Add bytes as possible field type

simw · Nov 13, 2023 · f55733c · f55733c
1 parent 812dc81
commit f55733c
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "pydantic_to_pyarrow"
-version = "0.1"
+version = "0.1.1"
 description = "Conversion from pydantic models to pyarrow schemas"
 authors = ["Simon Wicks <[email protected]>"]
 readme = "README.md"

diff --git a/src/pydantic_to_pyarrow/__init__.py b/src/pydantic_to_pyarrow/__init__.py
@@ -1,6 +1,6 @@
 from .schema import SchemaCreationError, get_pyarrow_schema
 
-__version__ = "0.1"
+__version__ = "0.1.1"
 
 __all__ = [
     "__version__",

diff --git a/src/pydantic_to_pyarrow/py.typed b/src/pydantic_to_pyarrow/py.typed
diff --git a/src/pydantic_to_pyarrow/schema.py b/src/pydantic_to_pyarrow/schema.py
@@ -17,6 +17,7 @@ class SchemaCreationError(Exception):
 
 FIELD_MAP = {
     str: pa.string(),
+    bytes: pa.binary(),
     bool: pa.bool_(),
     float: pa.float64(),
     datetime.date: pa.date32(),

diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -14,6 +14,7 @@
     NaiveDatetime,
     PositiveInt,
     StrictBool,
+    StrictBytes,
     StrictFloat,
     StrictInt,
     StrictStr,
@@ -51,20 +52,22 @@ class SimpleModel(BaseModel):
         b: bool
         c: int
         d: float
+        e: bytes
 
     expected = pa.schema(
         [
             pa.field("a", pa.string(), nullable=False),
             pa.field("b", pa.bool_(), nullable=False),
             pa.field("c", pa.int64(), nullable=False),
             pa.field("d", pa.float64(), nullable=False),
+            pa.field("e", pa.binary(), nullable=False),
         ]
     )
 
     actual = get_pyarrow_schema(SimpleModel)
     assert actual == expected
 
-    objs = [{"a": "a", "b": True, "c": 1, "d": 1.01}]
+    objs = [{"a": "a", "b": True, "c": 1, "d": 1.01, "e": b"e"}]
     new_schema, new_objs = _write_pq_and_read(objs, expected)
     assert new_schema == expected
     assert new_objs == objs
@@ -76,20 +79,22 @@ class SimpleModel(BaseModel):
         b: StrictBool
         c: StrictInt
         d: StrictFloat
+        e: StrictBytes
 
     expected = pa.schema(
         [
             pa.field("a", pa.string(), nullable=False),
             pa.field("b", pa.bool_(), nullable=False),
             pa.field("c", pa.int64(), nullable=False),
             pa.field("d", pa.float64(), nullable=False),
+            pa.field("e", pa.binary(), nullable=False),
         ]
     )
 
     actual = get_pyarrow_schema(SimpleModel)
     assert actual == expected
 
-    objs = [{"a": "a", "b": True, "c": 1, "d": 1.01}]
+    objs = [{"a": "a", "b": True, "c": 1, "d": 1.01, "e": b"e"}]
     new_schema, new_objs = _write_pq_and_read(objs, expected)
     assert new_schema == expected
     assert new_objs == objs