PrefectHQ / prefect

Prefect is a workflow orchestration framework for building resilient data pipelines in Python.
https://prefect.io
Apache License 2.0
15.31k stars 1.5k forks source link

handling of iterable field annotations in Block `schema_extra` #14282

Open WillRaphaelson opened 4 days ago

WillRaphaelson commented 4 days ago

in pydantic 1, field.type_ "unwrapped" the type inside the iterable generic alias (e.g. list) but that does not happen in v2

        if get_origin(field.annotation) in [Union, list]:

the below handles lists of block types, but not all iterable annotations

def schema_extra(schema: Dict[str, Any], model: Type["Block"]):
    """
    Customizes Pydantic's schema generation feature to add blocks related information.
    """
    schema["block_type_slug"] = model.get_block_type_slug()
    # Ensures args and code examples aren't included in the schema
    description = model.get_description()
    if description:
        schema["description"] = description
    else:
        # Prevent the description of the base class from being included in the schema
        schema.pop("description", None)

    # create a list of secret field names
    # secret fields include both top-level keys and dot-delimited nested secret keys
    # A wildcard (*) means that all fields under a given key are secret.
    # for example: ["x", "y", "z.*", "child.a"]
    # means the top-level keys "x" and "y", all keys under "z", and the key "a" of a block
    # nested under the "child" key are all secret. There is no limit to nesting.
    secrets = schema["secret_fields"] = []
    for name, field in model.model_fields.items():
        _collect_secret_fields(name, field.annotation, secrets)

    # create block schema references
    refs = schema["block_schema_references"] = {}
    for name, field in model.model_fields.items():
        if Block.is_block_class(field.annotation):
            refs[name] = field.annotation._to_block_schema_reference_dict()
        if get_origin(field.annotation) in [Union, list]:
            for type_ in get_args(field.annotation):
                if Block.is_block_class(type_):
                    if isinstance(refs.get(name), list):
                        refs[name].append(type_._to_block_schema_reference_dict())
                    elif isinstance(refs.get(name), dict):
                        refs[name] = [
                            refs[name],
                            type_._to_block_schema_reference_dict(),
                        ]
                    else:
                        refs[name] = type_._to_block_schema_reference_dict()