Skip to content

types

Module for handling type conversion.

make_nullable_array(schema)

Make an array nullable.

Source code in spark_instructor/utils/types.py 
31
32
33
def make_nullable_array(schema: ArrayType) -> ArrayType:
    """Make an array nullable."""
    return ArrayType(make_spark_schema_nullable(schema.elementType), True)

make_nullable_field(schema)

Make a field nullable.

Source code in spark_instructor/utils/types.py 
26
27
28
def make_nullable_field(schema: StructField) -> StructField:
    """Make a field nullable."""
    return StructField(schema.name, make_spark_schema_nullable(schema.dataType), nullable=True)

make_nullable_struct(schema)

Make a struct nullable.

Source code in spark_instructor/utils/types.py 
21
22
23
def make_nullable_struct(schema: StructType) -> StructType:
    """Make a struct nullable."""
    return StructType([make_spark_schema_nullable(field) for field in schema.fields])

make_spark_schema_nullable(schema)

Make a spark type nullable.

Source code in spark_instructor/utils/types.py 
36
37
38
39
40
41
42
43
44
45
def make_spark_schema_nullable(schema: T_Spark) -> T_Spark:
    """Make a spark type nullable."""
    if isinstance(schema, StructType):
        return make_nullable_struct(schema)
    elif isinstance(schema, StructField):
        return make_nullable_field(schema)
    elif isinstance(schema, ArrayType):
        return make_nullable_array(schema)
    else:
        return schema

pydantic_to_typeddict(pydantic_model, return_type, all_required=False)

Convert a pydantic model to a typed dict.

Source code in spark_instructor/utils/types.py 
48
49
50
51
52
53
54
def pydantic_to_typeddict(pydantic_model: BaseModel, return_type: Type[T], all_required: bool = False) -> T:
    """Convert a pydantic model to a typed dict."""
    result = pydantic_model.model_dump()
    if not all_required:
        result = _remove_none_from_dict(result)

    return return_type(**result)

typeddict_to_pydantic(typeddict_class)

Convert a TypedDict to a Pydantic model using TypeAdapter and core_schema.

Parameters:

Name Type Description Default
typeddict_class Type[Any]

The TypedDict class to convert.

 required

Returns:

Type Description
Type[BaseModel]

Type[BaseModel]: A Pydantic model class equivalent to the input TypedDict.
Example 
>>> from pydantic import BaseModel
>>> from typing import TypedDict
>>> class MyTypedDict(TypedDict):
...     name: str
...     age: int
>>> PydanticModel = typeddict_to_pydantic(MyTypedDict)
>>> isinstance(PydanticModel(name="John", age=30), BaseModel)
True
Source code in spark_instructor/utils/types.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def typeddict_to_pydantic(typeddict_class: Type[Any]) -> Type[BaseModel]:
    """Convert a TypedDict to a Pydantic model using TypeAdapter and core_schema.

    Args:
        typeddict_class (Type[Any]): The TypedDict class to convert.

    Returns:
        Type[BaseModel]: A Pydantic model class equivalent to the input TypedDict.

    Example:
        ```python

        >>> from pydantic import BaseModel
        >>> from typing import TypedDict
        >>> class MyTypedDict(TypedDict):
        ...     name: str
        ...     age: int
        >>> PydanticModel = typeddict_to_pydantic(MyTypedDict)
        >>> isinstance(PydanticModel(name="John", age=30), BaseModel)
        True

        ```
    """
    adapter = TypeAdapter(typeddict_class)
    core_schema = adapter.core_schema

    fields: Dict[str, tuple[Any, FieldInfo]] = {}
    _process_schema(core_schema, fields)

    return create_model(f"{typeddict_class.__name__}PD", __doc__=typeddict_class.__doc__, **fields)  # type: ignore