diff --git a/jambo/parser/enum_type_parser.py b/jambo/parser/enum_type_parser.py index 3d0fee3..e96aff5 100644 --- a/jambo/parser/enum_type_parser.py +++ b/jambo/parser/enum_type_parser.py @@ -36,6 +36,8 @@ class EnumTypeParser(GenericTypeParser): # Create a new Enum type dynamically enum_type = Enum(name, {str(value).upper(): value for value in enum_values}) # type: ignore + enum_type.__doc__ = properties.get("description") + parsed_properties = self.mappings_properties_builder(properties, **kwargs) if "default" in parsed_properties and parsed_properties["default"] is not None: 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_enum_type_parser.py b/tests/parser/test_enum_type_parser.py index 7665007..e180abf 100644 --- a/tests/parser/test_enum_type_parser.py +++ b/tests/parser/test_enum_type_parser.py @@ -49,6 +49,20 @@ class TestEnumTypeParser(TestCase): ) self.assertEqual(parsed_properties, {"default": None}) + def test_enum_type_parser_creates_enum_with_description(self): + parser = EnumTypeParser() + + schema = { + "description": "an enum", + "enum": ["value1"], + } + + parsed_type, parsed_properties = parser.from_properties_impl( + "TestEnum", + schema, + ) + self.assertEqual(parsed_type.__doc__, "an enum") + def test_enum_type_parser_creates_enum_with_default(self): parser = EnumTypeParser() 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"})