Register string datatype and update existing tests

Signed-off-by: Deepyaman Datta <[email protected]>
unionai-oss · Aug 19, 2024 · 55cf0f9 · 55cf0f9
1 parent b79ec31
commit 55cf0f9
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 12 deletions.
diff --git a/pandera/engines/ibis_engine.py b/pandera/engines/ibis_engine.py
@@ -118,3 +118,25 @@ class Float64(DataType, dtypes.Float64):
     """Semantic representation of a :class:`dt.Float64`."""
 
     type = dt.float64
+
+
+###############################################################################
+# nominal
+###############################################################################
+
+
+@Engine.register_dtype(
+    equivalents=[
+        str,
+        np.str_,
+        dtypes.String,
+        dtypes.String(),
+        dt.String,
+        dt.string,
+    ]
+)
+@immutable
+class String(DataType, dtypes.String):
+    """Semantic representation of a :class:`dt.String`."""
+
+    type = dt.string
diff --git a/pandera/engines/pandas_engine.py b/pandera/engines/pandas_engine.py
@@ -628,9 +628,9 @@ def __str__(self) -> str:
         return dtypes.Decimal.__str__(self)
 
 
-# ###############################################################################
-# # nominal
-# ###############################################################################
+###############################################################################
+# nominal
+###############################################################################
 
 
 @Engine.register_dtype(

diff --git a/tests/ibis/test_ibis_container.py b/tests/ibis/test_ibis_container.py
@@ -27,8 +27,8 @@ def t_schema_basic():
     """Basic Ibis table schema fixture."""
     return DataFrameSchema(
         {
-            # "string_col": Column(str),
-            "int_col": Column(int),
+            "string_col": Column(dt.String),
+            "int_col": Column(dt.Int64),
         }
     )
 
@@ -52,7 +52,7 @@ def test_required_columns():
     schema = DataFrameSchema(
         {
             "a": Column(dt.Int64, required=True),
-            "b": Column(dt.Int64, required=False),
+            "b": Column(dt.String, required=False),
         }
     )
     t = ibis.memtable(pd.DataFrame({"a": [1, 2, 3]}))

diff --git a/tests/ibis/test_ibis_model.py b/tests/ibis/test_ibis_model.py
@@ -11,7 +11,7 @@
 @pytest.fixture
 def t_model_basic():
     class BasicModel(DataFrameModel):
-        # string_col: str
+        string_col: str
         int_col: int
 
     return BasicModel
@@ -21,7 +21,7 @@ class BasicModel(DataFrameModel):
 def t_schema_basic():
     return DataFrameSchema(
         {
-            # "string_col": Column(dt.String),
+            "string_col": Column(dt.String),
             "int_col": Column(dt.Int64),
         }
     )
@@ -38,14 +38,14 @@ def test_model_schema_equivalency(
 
 def test_model_schema_equivalency_with_optional():
     class ModelWithOptional(DataFrameModel):
-        # string_col: Optional[str]
-        int_col: Optional[int]
+        string_col: Optional[str]
+        int_col: int
 
     schema = DataFrameSchema(
         name="ModelWithOptional",
         columns={
-            # "string_col": Column(dt.String, required=False),
-            "int_col": Column(dt.Int64, required=False),
+            "string_col": Column(dt.String, required=False),
+            "int_col": Column(dt.Int64),
         },
     )
     assert ModelWithOptional.to_schema() == schema