Partial expression evaluation, example of a builtin function

flow360/__init__.py

@@ -145,12 +145,12 @@
     SI_unit_system,
     imperial_unit_system,
 )
-from flow360.component.simulation.user_code import UserVariable
+from flow360.component.simulation.user_code.core.types import UserVariable
+from flow360.component.simulation.user_code.functions import math
+from flow360.component.simulation.user_code.variables import control, solution
 from flow360.component.simulation.user_defined_dynamics.user_defined_dynamics import (
     UserDefinedDynamic,
 )
-from flow360.component.simulation.variables import control_variables as control
-from flow360.component.simulation.variables import solution_variables as solution
 from flow360.component.surface_mesh_v2 import SurfaceMeshV2 as SurfaceMesh
 from flow360.component.volume_mesh import VolumeMeshV2 as VolumeMesh
 from flow360.environment import Env
@@ -277,4 +277,7 @@
     "Transformation",
     "WallRotation",
     "UserVariable",
+    "math",
+    "control",
+    "solution",
 ]

flow360/component/simulation/blueprint/__init__.py

@@ -6,7 +6,13 @@
     function_to_model,
 )
 
-from .core.function import Function
+from .core.function import FunctionNode
 from .core.types import Evaluable
 
-__all__ = ["Function", "Evaluable", "function_to_model", "model_to_function", "expr_to_model"]
+__all__ = [
+    "FunctionNode",
+    "Evaluable",
+    "function_to_model",
+    "model_to_function",
+    "expr_to_model",
+]

flow360/component/simulation/blueprint/core/__init__.py

@@ -2,30 +2,30 @@
 
 from .context import EvaluationContext, ReturnValue
 from .expressions import (
-    BinOp,
-    CallModel,
-    Constant,
-    Expression,
-    ExpressionType,
-    List,
-    ListComp,
-    Name,
-    RangeCall,
-    Subscript,
-    Tuple,
+    BinOpNode,
+    CallModelNode,
+    ConstantNode,
+    ExpressionNode,
+    ExpressionNodeType,
+    ListCompNode,
+    ListNode,
+    NameNode,
+    RangeCallNode,
+    SubscriptNode,
+    TupleNode,
 )
-from .function import Function
+from .function import FunctionNode
 from .generator import expr_to_code, model_to_function, stmt_to_code
 from .parser import function_to_model
 from .statements import (
-    Assign,
-    AugAssign,
-    ForLoop,
-    IfElse,
-    Return,
-    Statement,
-    StatementType,
-    TupleUnpack,
+    AssignNode,
+    AugAssignNode,
+    ForLoopNode,
+    IfElseNode,
+    ReturnNode,
+    StatementNode,
+    StatementNodeType,
+    TupleUnpackNode,
 )
 from .types import Evaluable, TargetSyntax
 
@@ -34,73 +34,73 @@ def _model_rebuild() -> None:
     """Update forward references in the correct order."""
     namespace = {
         # Expression types
-        "Name": Name,
-        "Constant": Constant,
-        "BinOp": BinOp,
-        "RangeCall": RangeCall,
-        "CallModel": CallModel,
-        "Tuple": Tuple,
-        "List": List,
-        "ListComp": ListComp,
-        "Subscript": Subscript,
-        "ExpressionType": ExpressionType,
+        "NameNode": NameNode,
+        "ConstantNode": ConstantNode,
+        "BinOpNode": BinOpNode,
+        "RangeCallNode": RangeCallNode,
+        "CallModelNode": CallModelNode,
+        "TupleNode": TupleNode,
+        "ListNode": ListNode,
+        "ListCompNode": ListCompNode,
+        "SubscriptNode": SubscriptNode,
+        "ExpressionNodeType": ExpressionNodeType,
         # Statement types
-        "Assign": Assign,
-        "AugAssign": AugAssign,
-        "IfElse": IfElse,
-        "ForLoop": ForLoop,
-        "Return": Return,
-        "TupleUnpack": TupleUnpack,
-        "StatementType": StatementType,
+        "AssignNode": AssignNode,
+        "AugAssignNode": AugAssignNode,
+        "IfElseNode": IfElseNode,
+        "ForLoopNode": ForLoopNode,
+        "ReturnNode": ReturnNode,
+        "TupleUnpackNode": TupleUnpackNode,
+        "StatementNodeType": StatementNodeType,
         # Function type
-        "Function": Function,
+        "FunctionNode": FunctionNode,
     }
 
     # First update expression classes that only depend on ExpressionType
-    BinOp.model_rebuild(_types_namespace=namespace)
-    RangeCall.model_rebuild(_types_namespace=namespace)
-    CallModel.model_rebuild(_types_namespace=namespace)
-    Tuple.model_rebuild(_types_namespace=namespace)
-    List.model_rebuild(_types_namespace=namespace)
-    ListComp.model_rebuild(_types_namespace=namespace)
-    Subscript.model_rebuild(_types_namespace=namespace)
+    BinOpNode.model_rebuild(_types_namespace=namespace)
+    RangeCallNode.model_rebuild(_types_namespace=namespace)
+    CallModelNode.model_rebuild(_types_namespace=namespace)
+    TupleNode.model_rebuild(_types_namespace=namespace)
+    ListNode.model_rebuild(_types_namespace=namespace)
+    ListCompNode.model_rebuild(_types_namespace=namespace)
+    SubscriptNode.model_rebuild(_types_namespace=namespace)
 
     # Then update statement classes that depend on both types
-    Assign.model_rebuild(_types_namespace=namespace)
-    AugAssign.model_rebuild(_types_namespace=namespace)
-    IfElse.model_rebuild(_types_namespace=namespace)
-    ForLoop.model_rebuild(_types_namespace=namespace)
-    Return.model_rebuild(_types_namespace=namespace)
-    TupleUnpack.model_rebuild(_types_namespace=namespace)
+    AssignNode.model_rebuild(_types_namespace=namespace)
+    AugAssignNode.model_rebuild(_types_namespace=namespace)
+    IfElseNode.model_rebuild(_types_namespace=namespace)
+    ForLoopNode.model_rebuild(_types_namespace=namespace)
+    ReturnNode.model_rebuild(_types_namespace=namespace)
+    TupleUnpackNode.model_rebuild(_types_namespace=namespace)
 
     # Finally update Function class
-    Function.model_rebuild(_types_namespace=namespace)
+    FunctionNode.model_rebuild(_types_namespace=namespace)
 
 
 # Update forward references
 _model_rebuild()
 
 
 __all__ = [
-    "Expression",
-    "Name",
-    "Constant",
-    "BinOp",
-    "RangeCall",
-    "CallModel",
-    "Tuple",
-    "List",
-    "ListComp",
-    "ExpressionType",
-    "Statement",
-    "Assign",
-    "AugAssign",
-    "IfElse",
-    "ForLoop",
-    "Return",
-    "TupleUnpack",
-    "StatementType",
-    "Function",
+    "ExpressionNode",
+    "NameNode",
+    "ConstantNode",
+    "BinOpNode",
+    "RangeCallNode",
+    "CallModelNode",
+    "TupleNode",
+    "ListNode",
+    "ListCompNode",
+    "ExpressionNodeType",
+    "StatementNode",
+    "AssignNode",
+    "AugAssignNode",
+    "IfElseNode",
+    "ForLoopNode",
+    "ReturnNode",
+    "TupleUnpackNode",
+    "StatementNodeType",
+    "FunctionNode",
     "EvaluationContext",
     "ReturnValue",
     "Evaluable",

flow360/component/simulation/blueprint/core/context.py

@@ -2,6 +2,8 @@
 
 from typing import Any, Optional
 
+import pydantic as pd
+
 from flow360.component.simulation.blueprint.core.resolver import CallableResolver
 
 
@@ -37,7 +39,9 @@ def __init__(
                 the context with.
         """
         self._values = initial_values or {}
+        self._data_models = {}
         self._resolver = resolver
+        self._aliases: dict[str, str] = {}
 
     def get(self, name: str, resolve: bool = True) -> Any:
         """
@@ -69,16 +73,34 @@ def get(self, name: str, resolve: bool = True) -> Any:
                 raise NameError(f"Name '{name}' is not defined") from err
         return self._values[name]
 
-    def set(self, name: str, value: Any) -> None:
+    def get_data_model(self, name: str) -> Optional[pd.BaseModel]:
+        """Get the Validation model for the given name."""
+        if name not in self._data_models:
+            return None
+        return self._data_models[name]
+
+    def set_alias(self, name, alias) -> None:
+        """Set alias used for code generation."""
+        self._aliases[name] = alias
+
+    def get_alias(self, name) -> Optional[str]:
+        """Get alias used for code generation."""
+        return self._aliases.get(name)
+
+    def set(self, name: str, value: Any, data_model: pd.BaseModel = None) -> None:
         """
         Assign a value to a name in the context.
 
         Args:
             name (str): The variable name to set.
             value (Any): The value to assign.
+            data_model (BaseModel, optional): The type of the associate with this entry
         """
         self._values[name] = value
 
+        if data_model:
+            self._data_models[name] = data_model
+
     def resolve(self, name):
         """
         Resolve a name using the provided resolver.