Address distinction between DataType and DataType() (#3857)

py-polars/polars/datatypes.py

@@ -21,7 +21,16 @@
 
 
 class DataType:
-    """Base class for all Polars data types"""
+    """
+    Base class for all Polars data types.
+    """
+
+    def __new__(cls, *args, **kwargs) -> "PolarsDataType":  # type: ignore
+        # this formulation allows for equivalent use of "pl.Type" and "pl.Type()", while
+        # still respecting types that take initialisation params (eg: Duration/Datetime)
+        if args or kwargs:
+            return super().__new__(cls)
+        return cls
 
     @classmethod
     def string_repr(cls) -> str:
@@ -182,7 +191,7 @@ def __eq__(self, other: Type[DataType]) -> bool:  # type: ignore
         # allow comparing object instances to class
         if type(other) is type and issubclass(other, Datetime):
             return True
-        if isinstance(other, Datetime):
+        elif isinstance(other, Datetime):
             return self.tu == other.tu and self.tz == other.tz
         else:
             return False
@@ -209,7 +218,7 @@ def __eq__(self, other: Type[DataType]) -> bool:  # type: ignore
         # allow comparing object instances to class
         if type(other) is type and issubclass(other, Duration):
             return True
-        if isinstance(other, Duration):
+        elif isinstance(other, Duration):
             return self.tu == other.tu
         else:
             return False

py-polars/polars/testing.py

@@ -709,6 +709,8 @@ def draw_frames(draw: Callable) -> Union[DataFrame, LazyFrame]:
                 )
                 dtypes_ = [draw(sampled_from(selectable_dtypes)) for _ in range(n)]
                 coldefs = columns(cols=n, dtype=dtypes_)
+            elif isinstance(cols, column):
+                coldefs = [cols]
             else:
                 coldefs = list(cols)  # type: ignore[arg-type]
 

py-polars/tests/test_datatypes.py

@@ -0,0 +1,28 @@
+import inspect
+
+import polars as pl
+import polars.datatypes as datatypes
+
+
+def test_dtype_init_equivalence() -> None:
+    # check "DataType.__new__" behaviour for all datatypes
+    all_datatypes = set(
+        dtype
+        for dtype in (getattr(datatypes, attr) for attr in dir(datatypes))
+        if inspect.isclass(dtype) and issubclass(dtype, datatypes.DataType)
+    )
+    for dtype in all_datatypes:
+        assert dtype == dtype()
+
+
+def test_dtype_temporal_units() -> None:
+    # check (in)equality behaviour of temporal types that take units
+    for tu in datatypes.DTYPE_TEMPORAL_UNITS:
+        assert pl.Datetime == pl.Datetime(tu)
+        assert pl.Duration == pl.Duration(tu)
+
+        assert pl.Datetime(tu) == pl.Datetime()  # type: ignore[operator]
+        assert pl.Duration(tu) == pl.Duration()  # type: ignore[operator]
+
+    assert pl.Datetime("ms") != pl.Datetime("ns")
+    assert pl.Duration("ns") != pl.Duration("us")