passing: types

Author: yaugenst-flex

tests/test_components/test_types.py

@@ -1,49 +1,33 @@
 """Tests type definitions."""
 
+from typing import Optional
+
 import numpy as np
-import pydantic.v1 as pydantic
 import pytest
+from pydantic import ValidationError
+from pydantic.fields import FieldInfo
 from tidy3d.components.base import Tidy3dBaseModel
-from tidy3d.components.types import ArrayLike, Complex, Tuple, constrained_array
-
-
-def _test_validate_array_like():
-    class S(Tidy3dBaseModel):
-        f: ArrayLike[float, 2]
-
-    _ = S(f=np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]))
-    with pytest.raises(pydantic.ValidationError):
-        _ = S(f=np.array([1.0, 2.0, 3.0]))
-
-    class MyClass(Tidy3dBaseModel):
-        f: constrained_array(ndim=3, shape=(1, 2, 3))
-
-    with pytest.raises(pydantic.ValidationError):
-        _ = MyClass(f=np.ones((2, 2, 3)))
-
-    with pytest.raises(pydantic.ValidationError):
-        _ = MyClass(f=np.ones((1, 2, 3, 4)))
+from tidy3d.components.types import ArrayConstraints, ArrayLike, Complex, array_alias
 
 
 def test_schemas():
     class S(Tidy3dBaseModel):
         f: ArrayLike
-        ca: constrained_array(ndim=1, dtype=complex)
+        ca: array_alias(ndim=1, dtype=complex)
         c: Complex
 
-    # TODO: unexpected behavior, if list with more than one element, it fails.
-    _ = S(f=[13], c=1 + 1j, ca=1 + 1j)
+    _ = S(f=[13], c=1 + 1j, ca=[1 + 1j])
     S.schema()
 
 
 def test_array_like():
     class MyClass(Tidy3dBaseModel):
         a: ArrayLike = None  # can be any array-like thing
-        b: constrained_array(ndim=2) = None  # must be 2D
-        c: constrained_array(dtype=float) = None  # must be float-like
-        d: constrained_array(ndim=1, dtype=complex) = None  # 1D complex
+        b: array_alias(ndim=2) = None  # must be 2D
+        c: array_alias(dtype=float) = None  # must be float-like
+        d: array_alias(ndim=1, dtype=complex) = None  # 1D complex
         e: ArrayLike
-        f: constrained_array(ndim=3, shape=(1, 2, 3)) = None  # must have certain shape
+        f: array_alias(ndim=3, shape=(1, 2, 3)) = None  # must have certain shape
 
     my_obj = MyClass(
         a=1.0 + 2j,
@@ -62,34 +46,311 @@ class MyClass(Tidy3dBaseModel):
     my_obj.json()
 
 
-def test_array_like_field_name():
+def test_array_like_field_metadata():
     class MyClass(Tidy3dBaseModel):
-        a: ArrayLike  # can be any array-like thing
-        b: constrained_array(ndim=2)  # must be 2D
-        c: constrained_array(dtype=float)  # must be float-like
-        d: constrained_array(ndim=1, dtype=complex)  # 1D complex
-        e: constrained_array(ndim=3, shape=(1, 2, 3))  # must have certain shape
-        f: ArrayLike = None
+        a: ArrayLike
+        b: array_alias(ndim=2)
+        c: array_alias(dtype=float)
+        d: array_alias(ndim=1, dtype=complex)
+        e: array_alias(ndim=3, shape=(1, 2, 3))
+        f: array_alias(strict=True, forbid_nan=False, scalar_to_1d=True)
+        g: array_alias(dtype=np.int32)
 
     fields = MyClass.model_fields
 
-    def correct_field_display(field_name, display_name):
-        """Make sure the field has the expected name."""
-        assert fields[field_name]._type_display() == display_name
+    def check_field_constraints(
+        field_name: str,
+        expected_dtype_name: Optional[str],
+        expected_ndim: Optional[int],
+        expected_shape: Optional[tuple[int, ...]],
+        expected_forbid_nan: bool,
+        expected_scalar_to_1d: bool,
+        expected_strict: bool,
+    ):
+        field_info: FieldInfo = fields[field_name]
 
-    correct_field_display("a", "ArrayLike")
-    correct_field_display("b", "ArrayLike[ndim=2]")
-    correct_field_display("c", "ArrayLike[dtype=float]")
-    correct_field_display("d", "ArrayLike[dtype=complex, ndim=1]")
-    correct_field_display("e", "ArrayLike[ndim=3, shape=(1, 2, 3)]")
-    correct_field_display("f", "Optional[ArrayLike]")
+        constraints = None
+        for meta_item in field_info.metadata:
+            if isinstance(meta_item, ArrayConstraints):
+                constraints = meta_item
+                break
+
+        assert (
+            constraints is not None
+        ), f"No ArrayConstraints found in metadata for field {field_name}. Metadata: {field_info.metadata}"
+
+        assert (
+            constraints.get("dtype") == expected_dtype_name
+        ), f"Field {field_name}: dtype mismatch. Expected {expected_dtype_name!r}, got {constraints.get('dtype')!r}"
+        assert (
+            constraints.get("ndim") == expected_ndim
+        ), f"Field {field_name}: ndim mismatch. Expected {expected_ndim}, got {constraints.get('ndim')}"
+        assert (
+            constraints.get("shape") == expected_shape
+        ), f"Field {field_name}: shape mismatch. Expected {expected_shape}, got {constraints.get('shape')}"
+        assert (
+            constraints.get("forbid_nan") == expected_forbid_nan
+        ), f"Field {field_name}: forbid_nan mismatch. Expected {expected_forbid_nan}, got {constraints.get('forbid_nan')}"
+        assert (
+            constraints.get("scalar_to_1d") == expected_scalar_to_1d
+        ), f"Field {field_name}: scalar_to_1d mismatch. Expected {expected_scalar_to_1d}, got {constraints.get('scalar_to_1d')}"
+        assert (
+            constraints.get("strict") == expected_strict
+        ), f"Field {field_name}: strict mismatch. Expected {expected_strict}, got {constraints.get('strict')}"
+
+    # Default values from array_alias: forbid_nan=True, scalar_to_1d=False, strict=False
+    check_field_constraints(
+        "a",
+        expected_dtype_name=None,
+        expected_ndim=None,
+        expected_shape=None,
+        expected_forbid_nan=True,
+        expected_scalar_to_1d=False,
+        expected_strict=False,
+    )
+
+    check_field_constraints(
+        "b",
+        expected_dtype_name=None,
+        expected_ndim=2,
+        expected_shape=None,
+        expected_forbid_nan=True,
+        expected_scalar_to_1d=False,
+        expected_strict=False,
+    )
+
+    # dtype=float -> float.__name__ is "float"
+    check_field_constraints(
+        "c",
+        expected_dtype_name="float",
+        expected_ndim=None,
+        expected_shape=None,
+        expected_forbid_nan=True,
+        expected_scalar_to_1d=False,
+        expected_strict=False,
+    )
+
+    # dtype=complex -> complex.__name__ is "complex"
+    check_field_constraints(
+        "d",
+        expected_dtype_name="complex",
+        expected_ndim=1,
+        expected_shape=None,
+        expected_forbid_nan=True,
+        expected_scalar_to_1d=False,
+        expected_strict=False,
+    )
+
+    check_field_constraints(
+        "e",
+        expected_dtype_name=None,
+        expected_ndim=3,
+        expected_shape=(1, 2, 3),
+        expected_forbid_nan=True,
+        expected_scalar_to_1d=False,
+        expected_strict=False,
+    )
+
+    # Test case for field 'f' with non-default values for flags
+    check_field_constraints(
+        "f",
+        expected_dtype_name=None,
+        expected_ndim=None,
+        expected_shape=None,
+        expected_forbid_nan=False,
+        expected_scalar_to_1d=True,
+        expected_strict=True,
+    )
+
+    # Test case for field 'g' with NumPy specific dtype
+    # dtype=np.int32 -> np.int32.__name__ is "int32"
+    check_field_constraints(
+        "g",
+        expected_dtype_name="int32",
+        expected_ndim=None,
+        expected_shape=None,
+        expected_forbid_nan=True,
+        expected_scalar_to_1d=False,
+        expected_strict=False,
+    )
 
 
 def test_hash():
     class MyClass(Tidy3dBaseModel):
         a: ArrayLike
-        b: constrained_array(ndim=1)
-        c: Tuple[ArrayLike, ...]
+        b: array_alias(ndim=1)
+        c: tuple[ArrayLike, ...]
 
     c = MyClass(a=[1.0], b=[2.0, 1.0], c=([2.0, 1.0]))
     hash(c.json())
+
+
+def test_array_like_validation_errors():
+    """Tests that appropriate ValidationErrors are raised for array constraints."""
+
+    # input that cannot be converted to a NumPy array at all (with specific dtype)
+    class ModelDtypeConversionFail(Tidy3dBaseModel):
+        a: array_alias(dtype=int)
+
+    with pytest.raises(ValidationError, match="cannot convert"):
+        ModelDtypeConversionFail(a="not an int")
+
+    # ndim mismatch
+    class ModelNdimMismatch(Tidy3dBaseModel):
+        a: array_alias(ndim=1)
+
+    with pytest.raises(ValidationError, match="expected"):
+        ModelNdimMismatch(a=[[1, 2], [3, 4]])
+
+    # ndim mismatch (scalar for ndim=1, scalar_to_1d=False by default)
+    class ModelNdimScalarDefault(Tidy3dBaseModel):
+        a: array_alias(ndim=1)
+
+    with pytest.raises(ValidationError, match="expected"):
+        ModelNdimScalarDefault(a=5)
+
+    # shape mismatch
+    class ModelShapeMismatch(Tidy3dBaseModel):
+        a: array_alias(shape=(2, 2))
+
+    with pytest.raises(ValidationError, match=r"expected shape"):
+        ModelShapeMismatch(a=[[1, 2, 3], [4, 5, 6]])
+
+    # forbid_nan=True (default) and array contains NaN
+    class ModelForbidNan(Tidy3dBaseModel):
+        a: array_alias(dtype=float)
+
+    with pytest.raises(ValidationError, match="array contains NaN"):
+        ModelForbidNan(a=[1.0, np.nan, 3.0])
+
+    # strict=True and a scalar is provided
+    class ModelStrictScalar(Tidy3dBaseModel):
+        a: array_alias(strict=True)
+
+    with pytest.raises(ValidationError, match="strict mode"):
+        ModelStrictScalar(a=10)
+
+    # input results in an array with dtype=object
+    class ModelObjectDtype(Tidy3dBaseModel):
+        a: ArrayLike
+
+    with pytest.raises(ValidationError, match=r"unsupported element type"):
+        ModelObjectDtype(a=[1, "string", object()])
+
+    # general conversion failure for an unhandled type
+    class ModelGeneralConversionFail(Tidy3dBaseModel):
+        a: ArrayLike
+
+    class UnconvertibleObject:
+        pass
+
+    with pytest.raises(ValidationError, match="unsupported element type"):
+        ModelGeneralConversionFail(a=UnconvertibleObject())
+
+    # _from_complex_dict receives a dict it doesn't understand, passes it to _coerce,
+    # which then fails because dict becomes an object array or direct conversion fails
+    class ModelComplexInvalidDict(Tidy3dBaseModel):
+        a: array_alias(dtype=complex)
+
+    with pytest.raises(ValidationError, match=r"cannot convert"):
+        ModelComplexInvalidDict(a={"real_part": 1, "imag_part": 2})
+
+    # scalar_to_1d=True with ndim=1 successfully converts scalar
+    class ModelScalarTo1DSuccess(Tidy3dBaseModel):
+        a: array_alias(ndim=1, scalar_to_1d=True)
+
+    obj_s21d = ModelScalarTo1DSuccess(a=5.0)
+    assert np.array_equal(obj_s21d.a, np.array([5.0]))
+    assert obj_s21d.a.ndim == 1
+
+    # scalar_to_1d=True but ndim is incompatible with 1D array (e.g. ndim=2)
+    class ModelScalarTo1DWrongNdim(Tidy3dBaseModel):
+        a: array_alias(ndim=2, scalar_to_1d=True, dtype=float)
+
+    with pytest.raises(ValidationError, match="expected"):
+        ModelScalarTo1DWrongNdim(a=5.0)
+
+    # strict=True takes precedence over scalar_to_1d=True if input is scalar
+    class ModelStrictAndScalarTo1D(Tidy3dBaseModel):
+        a: array_alias(strict=True, scalar_to_1d=True, dtype=float)
+
+    with pytest.raises(ValidationError, match="strict mode"):
+        ModelStrictAndScalarTo1D(a=5.0)
+
+    # allow NaN when forbid_nan=False
+    class ModelAllowNan(Tidy3dBaseModel):
+        a: array_alias(dtype=float, forbid_nan=False)
+
+    obj_allow_nan = ModelAllowNan(a=[1.0, np.nan])
+    assert np.array_equal(obj_allow_nan.a, np.array([1.0, np.nan]), equal_nan=True)
+
+    # strict=False (default) allows non-array if it can be coerced
+    class ModelStrictFalseCoercion(Tidy3dBaseModel):
+        a: array_alias(dtype=int, ndim=1)
+
+    # should pass because [1.0, 2.0] can be coerced to np.array([1,2]) of dtype int, ndim 1
+    obj_sf_coerce = ModelStrictFalseCoercion(a=[1.0, 2.0])
+    assert np.array_equal(obj_sf_coerce.a, np.array([1, 2]))
+    assert obj_sf_coerce.a.dtype == np.dtype(int)
+    assert obj_sf_coerce.a.ndim == 1
+
+    # scalar_to_1d=False (default), ndim=None, scalar input -> 0D array
+    class ModelScalarTo0D(Tidy3dBaseModel):
+        a: array_alias(scalar_to_1d=False)
+
+    obj_s0d = ModelScalarTo0D(a=10)
+    assert np.array_equal(obj_s0d.a, np.array(10))
+    assert obj_s0d.a.ndim == 0
+
+    # scalar_to_1d=True, ndim=None, scalar input -> 1D array
+    class ModelScalarTo1DNoNdim(Tidy3dBaseModel):
+        a: array_alias(scalar_to_1d=True)
+
+    obj_s1d_no_ndim = ModelScalarTo1DNoNdim(a=10)
+    assert np.array_equal(obj_s1d_no_ndim.a, np.array([10])), obj_s1d_no_ndim.a
+    assert obj_s1d_no_ndim.a.ndim == 1
+
+
+def test_complex_type():
+    """Tests the Complex type for parsing and serialization."""
+
+    class ComplexModel(Tidy3dBaseModel):
+        val: Complex
+
+    inputs = [
+        (1 + 2j, 1 + 2j),
+        ({"real": 3, "imag": -4}, 3 - 4j),
+        ({"real": 3.5, "imag": 0}, 3.5 + 0j),
+        (5, 5 + 0j),  # int
+        (6.7, 6.7 + 0j),  # float
+        (True, 1 + 0j),  # bool (subclass of int, numbers.Number)
+        (np.float32(2.5), 2.5 + 0j),  # numpy float
+        (np.int64(-3), -3 + 0j),  # numpy int
+        ([10, -2], 10 - 2j),  # list of two numbers
+        ((0.5, 1.5), 0.5 + 1.5j),  # tuple of two numbers
+    ]
+
+    class ObjWithComplexMethod:
+        def __complex__(self):
+            return -1 - 1j
+
+    class ObjWithComplexMethodNumeric:
+        def __init__(self, val):
+            self._val = val
+
+        def __complex__(self):
+            return self._val
+
+    inputs.append((ObjWithComplexMethod(), -1 - 1j))
+    inputs.append((ObjWithComplexMethodNumeric(3 + 7j), 3 + 7j))
+    inputs.append((ObjWithComplexMethodNumeric(5), 5 + 0j))
+
+    for input_val, expected_complex in inputs:
+        model = ComplexModel(val=input_val)
+        assert model.val == expected_complex, f"Input: {input_val}"
+        assert isinstance(model.val, complex), f"Input: {input_val}"
+
+        expected_json_val = {"real": expected_complex.real, "imag": expected_complex.imag}
+        assert (
+            model.model_dump(mode="json")["val"] == expected_json_val
+        ), f"Input for serialization: {input_val}"