mrd-python 2.0.0__py3-none-any.whl

mrd/_dtypes.py

@@ -0,0 +1,89 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import datetime
+from types import GenericAlias
+import sys
+
+if sys.version_info >= (3, 10):
+    from types import UnionType
+
+from typing import Any, Callable, Union, cast, get_args, get_origin
+import numpy as np
+from . import yardl_types as yardl
+
+
+def make_get_dtype_func(
+    dtype_map: dict[
+        Union[type, GenericAlias],
+        Union[np.dtype[Any], Callable[[tuple[type, ...]], np.dtype[Any]]],
+    ]
+) -> Callable[[Union[type, GenericAlias]], np.dtype[Any]]:
+    dtype_map[bool] = np.dtype(np.bool_)
+    dtype_map[yardl.Int8] = np.dtype(np.int8)
+    dtype_map[yardl.UInt8] = np.dtype(np.uint8)
+    dtype_map[yardl.Int16] = np.dtype(np.int16)
+    dtype_map[yardl.UInt16] = np.dtype(np.uint16)
+    dtype_map[yardl.Int32] = np.dtype(np.int32)
+    dtype_map[yardl.UInt32] = np.dtype(np.uint32)
+    dtype_map[yardl.Int64] = np.dtype(np.int64)
+    dtype_map[yardl.UInt64] = np.dtype(np.uint64)
+    dtype_map[yardl.Size] = np.dtype(np.uint64)
+    dtype_map[yardl.Float32] = np.dtype(np.float32)
+    dtype_map[yardl.Float64] = np.dtype(np.float64)
+    dtype_map[yardl.ComplexFloat] = np.dtype(np.complex64)
+    dtype_map[yardl.ComplexDouble] = np.dtype(np.complex128)
+    dtype_map[datetime.date] = np.dtype("datetime64[D]")
+    dtype_map[yardl.Time] = np.dtype("timedelta64[ns]")
+    dtype_map[yardl.DateTime] = np.dtype("datetime64[ns]")
+    dtype_map[str] = np.dtype(np.object_)
+
+    # Add the Python types to the dictionary too, but these may not be
+    # correct since they map to several dtypes
+    dtype_map[int] = np.dtype(np.int64)
+    dtype_map[float] = np.dtype(np.float64)
+    dtype_map[complex] = np.dtype(np.complex128)
+
+    dtype_map[list] = lambda type_args: np.dtype(np.object_)
+    dtype_map[dict] = lambda type_args: np.dtype(np.object_)
+    dtype_map[np.ndarray] = lambda type_args: np.dtype(np.object_)  # type: ignore
+
+    def get_dtype_impl(
+        dtype_map: dict[
+            Union[type, GenericAlias],
+            Union[np.dtype[Any], Callable[[tuple[type, ...]], np.dtype[Any]]],
+        ],
+        t: Union[type, GenericAlias],
+    ) -> np.dtype[Any]:
+        # Check dtype map for this type first
+        if (res := dtype_map.get(t, None)) is not None:
+            if callable(res):
+                raise RuntimeError(f"Generic type arguments not provided for {t}")
+            else:
+                return res
+
+        origin = get_origin(t)
+
+        if origin == Union or (
+            sys.version_info >= (3, 10) and isinstance(t, UnionType)
+        ):
+            return _get_union_dtype(get_args(t))
+
+        # Here, t is either invalid (no dtype registered)
+        # or t is a types.GenericAlias with type arguments specified
+        if origin is not None and (res := dtype_map.get(origin, None)) is not None:
+            if callable(res):
+                return res(get_args(t))
+
+        raise RuntimeError(f"Cannot find dtype for {t}")
+
+    def _get_union_dtype(args: tuple[type, ...]) -> np.dtype[Any]:
+        if len(args) == 2 and args[1] == cast(type, type(None)):
+            # This is an optional type
+            inner_type = get_dtype_impl(dtype_map, args[0])
+            return np.dtype(
+                [("has_value", np.bool_), ("value", inner_type)], align=True
+            )
+        return np.dtype(np.object_)
+
+    return lambda t: get_dtype_impl(dtype_map, t)