diff --git a/jambo/parser/object_type_parser.py b/jambo/parser/object_type_parser.py index 82ca587..5d611cd 100644 --- a/jambo/parser/object_type_parser.py +++ b/jambo/parser/object_type_parser.py @@ -22,6 +22,7 @@ class ObjectTypeParser(GenericTypeParser): name, properties.get("properties", {}), properties.get("required", []), + description=properties.get("description"), **kwargs, ) type_properties = self.mappings_properties_builder(properties, **kwargs) @@ -48,6 +49,7 @@ class ObjectTypeParser(GenericTypeParser): name: str, properties: dict[str, JSONSchema], required_keys: list[str], + description: str | None = None, **kwargs: Unpack[TypeParserOptions], ) -> type[BaseModel]: """ @@ -74,7 +76,9 @@ class ObjectTypeParser(GenericTypeParser): model_config = ConfigDict(validate_assignment=True) fields = cls._parse_properties(name, properties, required_keys, **kwargs) - model = create_model(name, __config__=model_config, **fields) # type: ignore + model = create_model( + name, __config__=model_config, __doc__=description, **fields + ) # type: ignore ref_cache[name] = model return model diff --git a/jambo/schema_converter.py b/jambo/schema_converter.py index 302b190..638c068 100644 --- a/jambo/schema_converter.py +++ b/jambo/schema_converter.py @@ -89,6 +89,7 @@ class SchemaConverter: schema["title"], schema.get("properties", {}), schema.get("required", []), + description=schema.get("description"), context=schema, ref_cache=ref_cache, required=True, diff --git a/tests/parser/test_object_type_parser.py b/tests/parser/test_object_type_parser.py index 975f85d..3bfd73a 100644 --- a/tests/parser/test_object_type_parser.py +++ b/tests/parser/test_object_type_parser.py @@ -24,6 +24,7 @@ class TestObjectTypeParser(TestCase): properties = { "type": "object", + "description": "obj desc", "properties": { "name": {"type": "string"}, "age": {"type": "integer"}, @@ -33,6 +34,7 @@ class TestObjectTypeParser(TestCase): Model, _args = parser.from_properties_impl( "placeholder", properties, ref_cache={} ) + self.assertEqual(Model.__doc__, "obj desc") obj = Model(name="name", age=10) diff --git a/tests/test_schema_converter.py b/tests/test_schema_converter.py index e9bcb5b..112d861 100644 --- a/tests/test_schema_converter.py +++ b/tests/test_schema_converter.py @@ -274,6 +274,7 @@ class TestSchemaConverter(TestCase): } model = self.converter.build_with_cache(schema) + self.assertEqual(model.__doc__, "A person") obj = model(address={"street": "123 Main St", "city": "Springfield"})