Merge pull request #1 from malinkinsa/0.0.2

Now, we correctly distinguish between root models and child models by checking if a model is only referenced but never itself a reference. And now child models correctly recognize nested models and process them as properties in Elasticsearch mappings.
malinkinsa · Mar 5, 2025 · 1e53b95 · 1e53b95
2 parents 8edbc31 + 901a5c3
commit 1e53b95
Show file tree

Hide file tree

Showing 6 changed files with 254 additions and 22 deletions.
diff --git a/pydantic2es/converters/dict2mapping.py b/pydantic2es/converters/dict2mapping.py
@@ -6,8 +6,7 @@
 def _create_mapping(data: dict, submodel_type: str, text_fields: List[str]) -> dict:
     mapping = {
         "mappings": {
-            "properties": {
-            }
+            "properties": {}
         }
     }
 
@@ -17,16 +16,14 @@ def _create_mapping(data: dict, submodel_type: str, text_fields: List[str]) -> d
                 "type": "text"
             }
 
+        elif isinstance(value, dict):
+            mapping['mappings']['properties'][key] = {
+                "type": submodel_type,
+                "properties": _create_mapping(value, submodel_type, text_fields)["mappings"]["properties"]
+            }
+
         else:
-            if isinstance(value, dict):
-                mapping['mappings']['properties'][key] = {
-                    "type": submodel_type,
-                    "properties": {}
-                }
-                for k, v in value.items():
-                    mapping['mappings']['properties'][key]['properties'][k] = get_mapping_value(v)
-            else:
-                mapping['mappings']['properties'][key] = get_mapping_value(value)
+            mapping['mappings']['properties'][key] = get_mapping_value(value)
 
     return mapping
 

diff --git a/pydantic2es/helpers/helpers.py b/pydantic2es/helpers/helpers.py
@@ -24,24 +24,22 @@ def get_mapping_value(key: str) -> dict:
     else:
         return {"type": mappings_map.get(key, None)}
 
-def _make_dict(converted_models: dict[dict[str, str]], parent_models: str, child_models_list: List[str]) \
-        -> dict:
-    data_dict = converted_models[parent_models]
+def _make_dict(converted_models: dict[str, dict[str, str]], parent_model: str, child_models_list: List[str]) -> dict:
+    data_dict = converted_models[parent_model].copy()
+
     for key, value in data_dict.items():
-        matching_child = next((child for child in child_models_list if child in value), None)
-        if matching_child:
-            data_dict[key] = converted_models[matching_child]
+        if value in child_models_list:
+            data_dict[key] = _make_dict(converted_models, value, child_models_list)
 
     return data_dict
 
+
 def struct_dict(converted_models: dict[dict[str, str]]) -> List[dict[str, str]] | dict[str, dict[str, str]]:
-    models_name = [converted_model_name for converted_model_name in converted_models]
     child_models_list = [
-        model
-        for model in models_name
+        value
         for sub_dict in converted_models.values()
         for value in sub_dict.values()
-        if model in value
+        if value in converted_models
     ]
 
     parent_models_list = [

diff --git a/setup.py b/setup.py
@@ -24,7 +24,7 @@ def readme():
 
 setup(
     name='pydantic-to-elastic',
-    version='0.0.1',
+    version='0.0.2',
     description='A simple CLI utility for converting Pydantic models to Elasticsearch mappings',
     license='MIT',
     long_description=readme(),

diff --git a/tests/__init__.py b/tests/__init__.py
diff --git a/tests/test_dict_to_mapping.py b/tests/test_dict_to_mapping.py
@@ -0,0 +1,201 @@
+import pytest
+
+from pydantic2es.converters.dict2mapping import dict_to_mapping
+
+
+@pytest.fixture
+def first_sample():
+    return {
+        'address': {
+            'city': 'str', 'street': 'str', 'zip_code': 'str'
+        },
+        'age': 'int',
+        'hobbies':
+            'list[str]',
+        'name': 'str'
+    }
+
+@pytest.fixture
+def second_sample():
+    return {
+        'user': {
+            'name': 'str',
+            'age': 'int',
+            'address': {
+                'city': 'str',
+                'street': 'str',
+                'zip_code': 'str',
+                'geo': {
+                    'latitude': 'float',
+                    'longitude': 'float'
+                }
+            },
+            'hobbies': 'list[str]'
+        }
+    }
+
+def test_dict_to_mapping(first_sample, second_sample):
+    first_mapping = dict_to_mapping(first_sample, 'nested', text_fields=[])
+    second_mapping = dict_to_mapping(second_sample, 'nested', text_fields=[])
+
+    expected_first_mapping = {
+        "mappings": {
+            "properties": {
+                "name": {
+                    "type": "keyword"
+                },
+                "age": {
+                    "type": "integer"
+                },
+                "address": {
+                    "type": "nested",
+                    "properties": {
+                        "street": {
+                            "type": "keyword"
+                        },
+                        "city": {
+                            "type": "keyword"
+                        },
+                        "zip_code": {
+                            "type": "keyword"
+                        }
+                    }
+                },
+                "hobbies": {
+                    "type": "keyword"
+                }
+            }
+        }
+    }
+
+    expected_second_mapping = {
+        "mappings": {
+            "properties": {
+                "user": {
+                    "type": "nested",
+                    "properties": {
+                        "name": {
+                            "type": "keyword"
+                        },
+                        "age": {
+                            "type": "integer"
+                        },
+                        "address": {
+                            "type": "nested",
+                            "properties": {
+                                "city": {
+                                    "type": "keyword"
+                                },
+                                "street": {
+                                    "type": "keyword"
+                                },
+                                "zip_code": {
+                                    "type": "keyword"
+                                },
+                                "geo": {
+                                    "type": "nested",
+                                    "properties": {
+                                        "latitude": {
+                                            "type": "float"
+                                        },
+                                        "longitude": {
+                                            "type": "float"
+                                        }
+                                    }
+                                }
+                            }
+                        },
+                        "hobbies": {
+                            "type": "keyword"
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    assert second_mapping == expected_second_mapping
+    assert first_mapping == expected_first_mapping
+
+def test_dict_to_mapping_with_test_and_object(first_sample, second_sample):
+    first_mapping = dict_to_mapping(first_sample, 'object', text_fields=['name'])
+    second_mapping = dict_to_mapping(second_sample, 'object', text_fields=['name'])
+
+    first_expected_mapping = {
+        "mappings": {
+            "properties": {
+                "name": {
+                    "type": "text"
+                },
+                "age": {
+                    "type": "integer"
+                },
+                "address": {
+                    "type": "object",
+                    "properties": {
+                        "street": {
+                            "type": "keyword"
+                        },
+                        "city": {
+                            "type": "keyword"
+                        },
+                        "zip_code": {
+                            "type": "keyword"
+                        }
+                    }
+                },
+                "hobbies": {
+                    "type": "keyword"
+                }
+            }
+        }
+    }
+
+    expected_second_mapping = {
+        "mappings": {
+            "properties": {
+                "user": {
+                    "type": "object",
+                    "properties": {
+                        "name": {
+                            "type": "text"
+                        },
+                        "age": {
+                            "type": "integer"
+                        },
+                        "address": {
+                            "type": "object",
+                            "properties": {
+                                "city": {
+                                    "type": "keyword"
+                                },
+                                "street": {
+                                    "type": "keyword"
+                                },
+                                "zip_code": {
+                                    "type": "keyword"
+                                },
+                                "geo": {
+                                    "type": "object",
+                                    "properties": {
+                                        "latitude": {
+                                            "type": "float"
+                                        },
+                                        "longitude": {
+                                            "type": "float"
+                                        }
+                                    }
+                                }
+                            }
+                        },
+                        "hobbies": {
+                            "type": "keyword"
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    assert second_mapping == expected_second_mapping
+    assert first_mapping == first_expected_mapping
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
@@ -0,0 +1,36 @@
+import pytest
+
+from pydantic2es.helpers.helpers import struct_dict
+
+
+@pytest.fixture
+def converted_models():
+    return {
+        'Address': {
+            'city': 'str',
+            'street': 'str',
+            'zip_code': 'str'
+        },
+        'User': {
+            'name': 'str',
+            'age': 'int',
+            'address': 'Address',
+            'hobbies': 'list[str]'
+        }
+    }
+
+
+def test_struct_dict_user_address(converted_models):
+    result = struct_dict(converted_models)
+
+    assert isinstance(result, dict)
+    assert isinstance(result["address"], dict)
+
+    assert 'User' not in result
+    assert "address" in result
+    assert "city" in result["address"]
+    assert "street" in result["address"]
+    assert "zip_code" in result["address"]
+    assert result["name"] == "str"
+    assert result["age"] == "int"
+    assert result["hobbies"] == "list[str]"