Merge pull request #78 from mikix/model-desc

feat: support object-level descriptions

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:

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

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,

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()
 

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)
 

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"})