Merge branch 'main' into feature/exclude-fields

simw · May 23, 2024 · 01b6c79 · 01b6c79
2 parents d2add23 + fe2cc1a
commit 01b6c79
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/src/pydantic_to_pyarrow/schema.py b/src/pydantic_to_pyarrow/schema.py
@@ -145,6 +145,16 @@ def _get_dict_type(
     )
 
 
+def _get_dict_type(
+    field_type: Type[Any], metadata: List[Any], allow_losing_tz: bool
+) -> pa.DataType:
+    key_type, value_type = get_args(field_type)
+    return pa.map_(
+        _get_pyarrow_type(key_type, metadata, allow_losing_tz=allow_losing_tz),
+        _get_pyarrow_type(value_type, metadata, allow_losing_tz=allow_losing_tz),
+    )
+
+
 FIELD_TYPES = {
     Literal: _get_literal_type,
     list: _get_list_type,

diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -531,7 +531,7 @@ class EnumModel(BaseModel):
     with pytest.raises(SchemaCreationError):
         get_pyarrow_schema(EnumModel)
 
-
+        
 def test_exclude_field_true() -> None:
     class SimpleModel(BaseModel):
         a: str
@@ -563,6 +563,8 @@ class SimpleModel(BaseModel):
     actual = get_pyarrow_schema(SimpleModel)
 
     assert actual == expected
+
+
 def test_dict() -> None:
     class DictModel(BaseModel):
         foo: Dict[str, int]