From 1d342bd7529380cdd57aea07aae789319cb8bf14 Mon Sep 17 00:00:00 2001
From: Joan Fontanals Martinez <joan.martinez@jina.ai>
Date: Sat, 29 Jul 2023 19:13:02 +0200
Subject: [PATCH] fix: add description and title to dynamic class

Signed-off-by: Joan Fontanals Martinez <joan.martinez@jina.ai>
---
 docarray/utils/create_dynamic_doc_class.py    | 19 +++++++++++++++++++
 .../util/test_create_dynamic_code_class.py    |  6 ++++++
 2 files changed, 25 insertions(+)

diff --git a/docarray/utils/create_dynamic_doc_class.py b/docarray/utils/create_dynamic_doc_class.py
index 6b85e97ee0..115257720b 100644
--- a/docarray/utils/create_dynamic_doc_class.py
+++ b/docarray/utils/create_dynamic_doc_class.py
@@ -6,6 +6,19 @@
 from docarray.utils._internal._typing import safe_issubclass
 
 
+RESERVED_KEYS = [
+    'type',
+    'anyOf',
+    '$ref',
+    'additionalProperties',
+    'allOf',
+    'items',
+    'definitions',
+    'properties',
+    'default',
+]
+
+
 def create_pure_python_type_model(model: Any) -> BaseDoc:
     """
     Take a Pydantic model and cast DocList fields into List fields.
@@ -239,5 +252,11 @@ class MyDoc(BaseDoc):
         )
 
     model = create_model(base_doc_name, __base__=BaseDoc, **fields)
+    model.__config__.title = schema.get('title', model.__config__.title)
+
+    for k in RESERVED_KEYS:
+        if k in schema:
+            schema.pop(k)
+    model.__config__.schema_extra = schema
     cached_models[base_doc_name] = model
     return model
diff --git a/tests/units/util/test_create_dynamic_code_class.py b/tests/units/util/test_create_dynamic_code_class.py
index 2e3243bfee..3e785d943e 100644
--- a/tests/units/util/test_create_dynamic_code_class.py
+++ b/tests/units/util/test_create_dynamic_code_class.py
@@ -243,6 +243,8 @@ class ResultTestDoc(BaseDoc):
 
 def test_create_with_field_info():
     class CustomDoc(BaseDoc):
+        """Here I have the description of the class"""
+
         a: str = Field(examples=['Example here'], another_extra='I am another extra')
 
     CustomDocCopy = create_pure_python_type_model(CustomDoc)
@@ -256,3 +258,7 @@ class CustomDoc(BaseDoc):
         new_custom_doc_model.schema().get('properties')['a']['another_extra']
         == 'I am another extra'
     )
+    assert (
+        new_custom_doc_model.schema().get('description')
+        == 'Here I have the description of the class'
+    )