Schema

schema

logger module-attribute

logger = logging.getLogger(__name__)

T module-attribute

T = TypeVar('T')

TModel module-attribute

TModel = TypeVar('TModel', bound=BaseModel)

_BaseModelType module-attribute

_BaseModelType = type(BaseModel)

CustomGenerateJsonSchema

CustomGenerateJsonSchema(*args, **kwargs)

Bases: GenerateJsonSchema

Custom JSON Schema generator to modify the way certain schemas are generated.

Source code in src/aind_behavior_services/schema/__init__.py

def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.nullable_as_oneof = kwargs.get("nullable_as_oneof", True)
    self.unions_as_oneof = kwargs.get("unions_as_oneof", True)
    self.render_x_enum_names = kwargs.get("render_x_enum_names", True)

nullable_as_oneof instance-attribute

nullable_as_oneof = kwargs.get('nullable_as_oneof', True)

unions_as_oneof instance-attribute

unions_as_oneof = kwargs.get('unions_as_oneof', True)

render_x_enum_names instance-attribute

render_x_enum_names = kwargs.get(
    "render_x_enum_names", True
)

nullable_schema

nullable_schema(schema)

Source code in src/aind_behavior_services/schema/__init__.py

def nullable_schema(self, schema: core_schema.NullableSchema) -> JsonSchemaValue:
    null_schema = {"type": "null"}
    inner_json_schema = self.generate_inner(schema["schema"])

    if inner_json_schema == null_schema:
        return null_schema
    else:
        if self.nullable_as_oneof:
            return self.get_flattened_oneof([inner_json_schema, null_schema])
        else:
            return super().get_flattened_anyof([inner_json_schema, null_schema])

get_flattened_oneof

get_flattened_oneof(schemas)

Source code in src/aind_behavior_services/schema/__init__.py

def get_flattened_oneof(self, schemas: list[JsonSchemaValue]) -> JsonSchemaValue:
    members = []
    for schema in schemas:
        if len(schema) == 1 and "oneOf" in schema:
            members.extend(schema["oneOf"])
        else:
            members.append(schema)
    members = _deduplicate_schemas(members)
    if len(members) == 1:
        return members[0]
    return {"oneOf": members}

enum_schema

enum_schema(schema)

Generates a JSON schema that matches an Enum value.

PARAMETER	DESCRIPTION
schema	The core schema. TYPE: EnumSchema

RETURNS	DESCRIPTION
JsonSchemaValue	The generated JSON schema.

Source code in src/aind_behavior_services/schema/__init__.py

def enum_schema(self, schema: core_schema.EnumSchema) -> JsonSchemaValue:
    """Generates a JSON schema that matches an Enum value.

    Args:
        schema: The core schema.

    Returns:
        The generated JSON schema.
    """
    enum_type = schema["cls"]
    description = None if not enum_type.__doc__ else inspect.cleandoc(enum_type.__doc__)
    if (
        description == "An enumeration."
    ):  # This is the default value provided by enum.EnumMeta.__new__; don't use it
        description = None
    result: dict[str, Any] = {"title": enum_type.__name__, "description": description}
    result = {k: v for k, v in result.items() if v is not None}

    expected = [to_jsonable_python(v.value) for v in schema["members"]]

    result["enum"] = expected
    if len(expected) == 1:
        result["const"] = expected[0]

    types = {type(e) for e in expected}
    if isinstance(enum_type, str) or types == {str}:
        result["type"] = "string"
    elif isinstance(enum_type, int) or types == {int}:
        result["type"] = "integer"
    elif isinstance(enum_type, float) or types == {float}:
        result["type"] = "numeric"
    elif types == {bool}:
        result["type"] = "boolean"
    elif types == {list}:
        result["type"] = "array"

    _type = result.get("type", None)
    if (self.render_x_enum_names) and (_type != "string"):
        result["x-enumNames"] = [screaming_snake_case_to_pascal_case(v.name) for v in schema["members"]]

    return result

literal_schema

literal_schema(schema)

Generates a JSON schema that matches a literal value.

PARAMETER	DESCRIPTION
schema	The core schema. TYPE: LiteralSchema

RETURNS	DESCRIPTION
JsonSchemaValue	The generated JSON schema.

Source code in src/aind_behavior_services/schema/__init__.py

def literal_schema(self, schema: core_schema.LiteralSchema) -> JsonSchemaValue:
    """Generates a JSON schema that matches a literal value.

    Args:
        schema: The core schema.

    Returns:
        The generated JSON schema.
    """
    expected = [v.value if isinstance(v, Enum) else v for v in schema["expected"]]
    # jsonify the expected values
    expected = [to_jsonable_python(v) for v in expected]

    types = {type(e) for e in expected}

    if len(expected) == 1:
        if isinstance(expected[0], str):
            return {"const": expected[0], "type": "string"}
        elif isinstance(expected[0], bool):
            return {"const": expected[0], "type": "boolean"}
        elif isinstance(expected[0], int):
            return {"const": expected[0], "type": "integer"}
        elif isinstance(expected[0], float):
            return {"const": expected[0], "type": "number"}
        elif isinstance(expected[0], list):
            return {"const": expected[0], "type": "array"}
        elif expected[0] is None:
            return {"const": expected[0], "type": "null"}
        else:
            return {"const": expected[0]}

    if types == {str}:
        return {"enum": expected, "type": "string"}
    elif types == {bool}:
        return {"enum": expected, "type": "boolean"}
    elif types == {int}:
        return {"enum": expected, "type": "integer"}
    elif types == {float}:
        return {"enum": expected, "type": "number"}
    elif types == {list}:
        return {"enum": expected, "type": "array"}
    # there is not None case because if it's mixed it hits the final `else`
    # if it's a single Literal[None] then it becomes a `const` schema above
    else:
        return {"enum": expected}

generate_inner

generate_inner(schema)

Source code in src/aind_behavior_services/schema/__init__.py

def generate_inner(self, schema: core_schema.CoreSchema) -> JsonSchemaValue:
    defs_before = set(self.definitions.keys())
    result = super().generate_inner(schema)
    new_refs = set(self.definitions.keys()) - defs_before

    if new_refs:
        cls = schema.get("cls")
        typename: Optional[str] = None
        if cls is not None and not (isinstance(cls, type) and issubclass(cls, BaseModel)):
            typename = cls.__dict__.get("__sgen_typename__")
        if typename is not None:
            for ref in new_refs:
                self.definitions[ref]["x-sgen-typename"] = typename

    return result

union_schema

union_schema(schema)

Generates a JSON schema that matches a schema that allows values matching any of the given schemas.

PARAMETER	DESCRIPTION
schema	The core schema. TYPE: UnionSchema

RETURNS	DESCRIPTION
JsonSchemaValue	The generated JSON schema.

Source code in src/aind_behavior_services/schema/__init__.py

def union_schema(self, schema: core_schema.UnionSchema) -> JsonSchemaValue:
    """Generates a JSON schema that matches a schema that allows values matching any of the given schemas.

    Args:
        schema: The core schema.

    Returns:
        The generated JSON schema.
    """
    generated: list[JsonSchemaValue] = []

    choices = schema["choices"]
    for choice in choices:
        # choice will be a tuple if an explicit label was provided
        choice_schema = choice[0] if isinstance(choice, tuple) else choice
        try:
            generated.append(self.generate_inner(choice_schema))
        except PydanticOmit:
            continue
        except PydanticInvalidForJsonSchema as exc:
            self.emit_warning("skipped-choice", exc.message)
    if len(generated) == 1:
        return generated[0]
    if self.unions_as_oneof is True:
        return self.get_flattened_oneof(generated)
    else:
        return self.get_flattened_anyof(generated)

BonsaiSgenSerializers

Bases: Enum

NONE class-attribute instance-attribute

NONE = 'None'

JSON class-attribute instance-attribute

JSON = 'json'

YAML class-attribute instance-attribute

YAML = 'yaml'

_SgenBaseModelMeta

Bases: _BaseModelType

Metaclass that strips x-sgen-typename from subclasses of sgen_typename-decorated models. Subclasses must re-decorate to opt in.

_SgenTypenameAnnotation

_SgenTypenameAnnotation(typename)

Annotation marker for frozen types (e.g., TypeAliasType) that injects x-sgen-typename into the type's $defs entry via pydantic_js_updates.

Source code in src/aind_behavior_services/schema/__init__.py

def __init__(self, typename: str) -> None:
    self._typename = typename

SgenNamespace

SgenNamespace(namespace)

Source code in src/aind_behavior_services/schema/__init__.py

def __init__(self, namespace: str):
    self._namespace = namespace

namespace property

namespace

sgen_typename

sgen_typename(*, typename=None)

Source code in src/aind_behavior_services/schema/__init__.py

def sgen_typename(self, *, typename: str | None = None) -> Callable[[Type[T]], Type[T]]:
    return sgen_typename(typename=typename, namespace=self._namespace)

screaming_snake_case_to_pascal_case

screaming_snake_case_to_pascal_case(s)

Converts a SCREAMING_SNAKE_CASE string to PascalCase.

PARAMETER	DESCRIPTION
s	The SCREAMING_SNAKE_CASE string to be converted. TYPE: str

RETURNS	DESCRIPTION
str	The PascalCase string. TYPE: str

Source code in src/aind_behavior_services/utils.py

def screaming_snake_case_to_pascal_case(s: str) -> str:
    """
    Converts a SCREAMING_SNAKE_CASE string to PascalCase.

    Args:
        s (str): The SCREAMING_SNAKE_CASE string to be converted.

    Returns:
        str: The PascalCase string.
    """
    words = s.split("_")
    return "".join(word.capitalize() for word in words)

snake_to_pascal_case

snake_to_pascal_case(s)

Converts a snake_case string to PascalCase.

PARAMETER	DESCRIPTION
s	The snake_case string to be converted. TYPE: str

RETURNS	DESCRIPTION
str	The PascalCase string. TYPE: str

Source code in src/aind_behavior_services/utils.py

def snake_to_pascal_case(s: str) -> str:
    """
    Converts a snake_case string to PascalCase.

    Args:
        s (str): The snake_case string to be converted.

    Returns:
        str: The PascalCase string.
    """
    return "".join(map(capwords, s.split("_")))

export_schema

export_schema(
    model,
    schema_generator=CustomGenerateJsonSchema,
    mode="serialization",
    remove_root=True,
)

Export the schema of a model to a json file

Source code in src/aind_behavior_services/schema/__init__.py

def export_schema(
    model: Type[BaseModel],
    schema_generator: Type[GenerateJsonSchema] = CustomGenerateJsonSchema,
    mode: JsonSchemaMode = "serialization",
    remove_root: bool = True,
):
    """Export the schema of a model to a json file"""
    _model = model.model_json_schema(schema_generator=schema_generator, mode=mode)
    if remove_root:
        for to_remove in ["title", "description", "properties", "required", "type", "oneOf"]:
            _model.pop(to_remove, None)
    return json.dumps(_model, indent=2)

bonsai_sgen

bonsai_sgen(
    schema_path,
    output_path,
    namespace=None,
    root_element=None,
    serializer=None,
)

Runs Bonsai.SGen to generate a Bonsai-compatible schema from a json-schema model For more information run bonsai.sgen --help in the command line.

RETURNS	DESCRIPTION
CompletedProcess	The result of running the command. TYPE: CompletedProcess

Args: schema_path (PathLike): Target Json Schema file output_path (PathLike): Specifies the name of the file containing the generated code. namespace (Optional[str], optional): Specifies the namespace to use for all generated serialization classes. Defaults to DataSchema. root_element (Optional[str], optional): Specifies the name of the class used to represent the schema root element. If None, it will use the json schema root element. Defaults to None. serializer (Optional[List[BonsaiSgenSerializers]], optional): Specifies the serializer data annotations to include in the generated classes. Defaults to None.

Source code in src/aind_behavior_services/schema/__init__.py

def bonsai_sgen(
    schema_path: PathLike,
    output_path: PathLike,
    namespace: Optional[str] = None,
    root_element: Optional[str] = None,
    serializer: Optional[List[BonsaiSgenSerializers]] = None,
) -> CompletedProcess:
    """Runs Bonsai.SGen to generate a Bonsai-compatible schema from a json-schema model
    For more information run `bonsai.sgen --help` in the command line.

    Returns:
        CompletedProcess: The result of running the command.
    Args:
        schema_path (PathLike): Target Json Schema file
        output_path (PathLike): Specifies the name of the
          file containing the generated code.
        namespace (Optional[str], optional): Specifies the
          namespace to use for all generated serialization
          classes. Defaults to DataSchema.
        root_element (Optional[str], optional):  Specifies the
          name of the class used to represent the schema root element.
          If None, it will use the json schema root element. Defaults to None.
        serializer (Optional[List[BonsaiSgenSerializers]], optional):
          Specifies the serializer data annotations to include in the generated classes.
          Defaults to None.
    """

    if serializer is None:
        serializer = [BonsaiSgenSerializers.JSON]

    _restore_cmd = run("dotnet tool restore", shell=True, check=True, capture_output=True)
    try:
        _restore_cmd.check_returncode()
    except CalledProcessError as e:
        print(f"Error occurred while restoring tools: {e}")
        print(
            "Ensure you have the Bonsai.Sgen tool installed locally. "
            "See https://github.com/bonsai-rx/sgen?tab=readme-ov-file#getting-started for instructions."
        )
        raise

    version = _check_bonsai_sgen_version()
    if version < Version.parse("0.6.0"):
        raise RuntimeError("Version of Bonsai.Sgen must be at least 0.6.0, found: " + str(version))

    cmd_string = (
        f'dotnet tool run bonsai.sgen "{schema_path}" -o "{Path(output_path).parent}" -n {Path(output_path).name}'
    )
    cmd_string += f" --namespace {namespace}" if namespace is not None else ""
    cmd_string += f" --root {root_element}" if root_element is not None else ""

    if len(serializer) == 0 or BonsaiSgenSerializers.NONE in serializer:
        cmd_string += " --serializer none"
    else:
        cmd_string += " --serializer"
        cmd_string += " ".join([f" {sr.value}" for sr in serializer])
    return run(cmd_string, shell=True, check=True)

_check_bonsai_sgen_version

_check_bonsai_sgen_version()

Check the version of the Bonsai.SGen tool.

Source code in src/aind_behavior_services/schema/__init__.py

def _check_bonsai_sgen_version() -> Version:
    """Check the version of the Bonsai.SGen tool."""
    result = run("dotnet tool run bonsai.sgen --version", shell=True, check=True, capture_output=True)
    version_str = result.stdout.strip()
    return Version.parse(version_str)

convert_pydantic_to_bonsai

convert_pydantic_to_bonsai(
    model,
    *,
    model_name=None,
    json_schema_output_dir=Path("./src/DataSchemas/"),
    cs_output_dir=Path("./src/Extensions/"),
    cs_namespace="DataSchema",
    cs_serializer=None,
    json_schema_export_kwargs=None,
    root_element=None,
)

Source code in src/aind_behavior_services/schema/__init__.py

def convert_pydantic_to_bonsai(
    model: Type[BaseModel],
    *,
    model_name: Optional[str] = None,
    json_schema_output_dir: PathLike = Path("./src/DataSchemas/"),
    cs_output_dir: Optional[PathLike] = Path("./src/Extensions/"),
    cs_namespace: str = "DataSchema",
    cs_serializer: Optional[List[BonsaiSgenSerializers]] = None,
    json_schema_export_kwargs: Optional[Dict[str, Any]] = None,
    root_element: Optional[str] = None,
) -> Optional[CompletedProcess]:
    def _write_json(schema_path: PathLike, output_model_name: str, model: Type[BaseModel], **extra_kwargs) -> None:
        with open(os.path.join(schema_path, f"{output_model_name}.json"), "w", encoding="utf-8") as f:
            json_model = export_schema(model, **extra_kwargs)
            f.write(json_model)

    _model_name = model_name or model.__name__
    _write_json(json_schema_output_dir, _model_name, model, **(json_schema_export_kwargs or {}))

    if cs_output_dir is not None:
        cmd_return = bonsai_sgen(
            schema_path=Path(os.path.join(json_schema_output_dir, f"{_model_name}.json")),
            output_path=Path(os.path.join(cs_output_dir, f"{snake_to_pascal_case(_model_name)}.Generated.cs")),
            namespace=cs_namespace,
            serializer=cs_serializer,
            root_element=root_element,
        )
        return cmd_return

    return None

sgen_typename

sgen_typename(*, typename=None, namespace=None)

Class decorator to add an x-sgen-typename property to the model's JSON schema, which Bonsai.SGen uses to determine the typename for the generated class.

Source code in src/aind_behavior_services/schema/__init__.py

def sgen_typename(*, typename: Optional[str] = None, namespace: str | None = None) -> Callable[[Type[T]], Type[T]]:
    """Class decorator to add an ``x-sgen-typename`` property to the model's JSON schema, which Bonsai.SGen uses to determine the typename for the generated class."""

    # Why do we need 3 different approaches you might ask?
    # 1. For regular BaseModel subclasses, we create_model subclass because we want the strip the x-sgen-typename from any subclasses by default (to avoid unintended inheritance) and require explicit re-decoration opt-in. This is handled by the _SgenBaseModelMeta metaclass.
    # 2. For all other classes, we can just set the __sgen_typename__ attribute directly, and the custom JSON schema generator will pick it up and add it to the generated schema. This is the simplest case.
    # 3. 2. For frozen types (e.g., TypeAliasType), we can't modify using 2. so we wrap the type in an Annotated with a custom annotation that injects the x-sgen-typename via pydantic_js_updates. Handled by the _SgenTypenameAnnotation class.

    def decorator(cls: Type[T]) -> Type[T]:
        _typename = typename or cls.__name__
        _typename = f"{namespace}.{_typename}" if namespace else _typename
        if isinstance(cls, type) and issubclass(cls, BaseModel):
            existing = getattr(cls, "model_config", ConfigDict())
            raw_extra = existing.get("json_schema_extra")
            new_extra: Dict[str, Any] = dict(raw_extra) if isinstance(raw_extra, dict) else {}
            new_extra["x-sgen-typename"] = _typename
            new_config = cast(ConfigDict, {**existing, "json_schema_extra": new_extra})
            result = create_model(
                cls.__name__,
                __base__=cls,
                __config__=new_config,
                __module__=cls.__module__,
                __cls_kwargs__={"metaclass": _SgenBaseModelMeta},
                __doc__=cls.__doc__,
            )
            result.__qualname__ = cls.__qualname__  # type: ignore[attr-defined]
        else:
            result = cls  # type: ignore[assignment]
        try:
            setattr(result, "__sgen_typename__", _typename)
        except AttributeError:
            # Frozen object (e.g., TypeAliasType); wrap in Annotated with a marker
            # that injects x-sgen-typename into the $defs entry via pydantic_js_functions.
            result = Annotated[result, _SgenTypenameAnnotation(_typename)]  # type: ignore[assignment]
        return cast(Type[T], result)

    return decorator