Made things actually serializable

tidy3d/components/expressions/expressions.py

@@ -1,24 +1,25 @@
-import numbers
+import pydantic.v1 as pd
 
+from ..types import TYPE_TAG_STR
 from .abstract_base import AbstractExpression
 from .operators import (
     Absolute,
     Division,
     FloorDiv,
     Imaginary,
     Modulo,
-    Operator,
     Power,
     Product,
     Real,
     Sum,
 )
+from .types import ExpressionType, NumberType, OperatorType
 
 
 class Expression(AbstractExpression):
     @staticmethod
     def _ensure_expression(value):
-        if isinstance(value, numbers.Number):
+        if isinstance(value, NumberType):
             return Constant(value=value)
         return value
 
@@ -87,16 +88,19 @@ def imag(self):
 
 
 class Constant(Expression):
-    value: numbers.Number
+    value: NumberType = pd.Field()
+
+    class Config:
+        smart_union = True  # to preserve input type
 
     def evaluate(self, _):
         return self.value
 
 
 class CompoundExpression(Expression):
-    left: Expression
-    right: Expression
-    operator: Operator
+    left: ExpressionType = pd.Field(discriminator=TYPE_TAG_STR)
+    right: ExpressionType = pd.Field(discriminator=TYPE_TAG_STR)
+    operator: OperatorType = pd.Field(discriminator=TYPE_TAG_STR)
 
     def evaluate(self, data):
         return self.operator(self.left(data), self.right(data))

tidy3d/components/expressions/types.py

@@ -0,0 +1,10 @@
+from typing import ForwardRef, Union
+
+from ..types import ArrayLike, Complex
+from .operators import Absolute, Division, FloorDiv, Imaginary, Modulo, Power, Product, Real, Sum
+
+OperatorType = Union[Sum, Product, Division, Absolute, FloorDiv, Modulo, Power, Real, Imaginary]
+ExpressionType = Union[
+    ForwardRef("Constant"), ForwardRef("CompoundExpression")
+]  # forward refs because `CompoundExpression` has to reference itself
+NumberType = Union[int, float, Complex, ArrayLike]