OpenFisca-Core-test 46.0.0__py3-none-any.whl

openfisca_core/tracers/trace_node.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import dataclasses
+
+from openfisca_core import types as t
+
+
+@dataclasses.dataclass
+class TraceNode:
+    """A node in the tracing tree."""
+
+    #: The name of the node.
+    name: str
+
+    #: The period of the node.
+    period: t.PeriodInt | t.Period
+
+    #: The parent of the node.
+    parent: None | t.TraceNode = None
+
+    #: The children of the node.
+    children: list[t.TraceNode] = dataclasses.field(default_factory=list)
+
+    #: The parameters of the node.
+    parameters: list[t.TraceNode] = dataclasses.field(default_factory=list)
+
+    #: The value of the node.
+    value: None | t.VarArray = None
+
+    #: The start time of the node.
+    start: t.Time = 0.0
+
+    #: The end time of the node.
+    end: t.Time = 0.0
+
+    def calculation_time(self, round_: bool = True) -> t.Time:
+        """Calculate the time spent in the node.
+
+        Args:
+            round_: Whether to round the result.
+
+        Returns:
+            float: The time spent in the node.
+
+        Examples:
+            >>> from openfisca_core import tracers
+
+            >>> node = tracers.TraceNode("variable", 2020)
+            >>> node.start = 1.123122313
+            >>> node.end = 1.12312313123
+
+            >>> node.calculation_time()
+            8.182e-07
+
+            >>> node.calculation_time(round_=False)
+            8.182299999770493e-07
+
+        """
+        result = self.end - self.start
+
+        if round_:
+            return self.round(result)
+
+        return result
+
+    def formula_time(self) -> t.Time:
+        """Calculate the time spent on the formula.
+
+        Returns:
+            float: The time spent on the formula.
+
+        Examples:
+            >>> from openfisca_core import tracers
+
+            >>> node = tracers.TraceNode("variable", 2020)
+            >>> node.start = 1.123122313 * 11
+            >>> node.end = 1.12312313123 * 11
+            >>> child = tracers.TraceNode("variable", 2020)
+            >>> child.start = 1.123122313
+            >>> child.end = 1.12312313123
+
+            >>> for i in range(10):
+            ...     node.children = [child, *node.children]
+
+            >>> node.formula_time()
+            8.182e-07
+
+        """
+        children_calculation_time = sum(
+            child.calculation_time(round_=False) for child in self.children
+        )
+
+        result = +self.calculation_time(round_=False) - children_calculation_time
+
+        return self.round(result)
+
+    def append_child(self, node: t.TraceNode) -> None:
+        """Append a child to the node."""
+        self.children.append(node)
+
+    @staticmethod
+    def round(time: t.Time) -> t.Time:
+        """Keep only 4 significant figures.
+
+        Args:
+            time: The time to round.
+
+        Returns:
+            float: The rounded time.
+
+        Examples:
+            >>> from openfisca_core import tracers
+
+            >>> tracers.TraceNode.round(0.000123456789)
+            0.0001235
+
+        """
+        return float(f"{time:.4g}")
+
+
+__all__ = ["TraceNode"]

openfisca_core/tracers/tracing_parameter_node_at_instant.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import typing
+
+import numpy
+
+from openfisca_core import parameters
+
+ParameterNode = (
+    parameters.VectorialParameterNodeAtInstant | parameters.ParameterNodeAtInstant
+)
+
+if typing.TYPE_CHECKING:
+    from numpy.typing import ArrayLike
+
+    from openfisca_core import tracers
+
+    Child = ParameterNode | ArrayLike
+
+
+class TracingParameterNodeAtInstant:
+    def __init__(
+        self,
+        parameter_node_at_instant: ParameterNode,
+        tracer: tracers.FullTracer,
+    ) -> None:
+        self.parameter_node_at_instant = parameter_node_at_instant
+        self.tracer = tracer
+
+    def __getattr__(
+        self,
+        key: str,
+    ) -> TracingParameterNodeAtInstant | Child:
+        child = getattr(self.parameter_node_at_instant, key)
+        return self.get_traced_child(child, key)
+
+    def __contains__(self, key) -> bool:
+        return key in self.parameter_node_at_instant
+
+    def __iter__(self):
+        return iter(self.parameter_node_at_instant)
+
+    def __getitem__(
+        self,
+        key: str | ArrayLike,
+    ) -> TracingParameterNodeAtInstant | Child:
+        child = self.parameter_node_at_instant[key]
+        return self.get_traced_child(child, key)
+
+    def get_traced_child(
+        self,
+        child: Child,
+        key: str | ArrayLike,
+    ) -> TracingParameterNodeAtInstant | Child:
+        period = self.parameter_node_at_instant._instant_str
+
+        if isinstance(
+            child,
+            (
+                parameters.ParameterNodeAtInstant,
+                parameters.VectorialParameterNodeAtInstant,
+            ),
+        ):
+            return TracingParameterNodeAtInstant(child, self.tracer)
+
+        if not isinstance(key, str) or isinstance(
+            self.parameter_node_at_instant,
+            parameters.VectorialParameterNodeAtInstant,
+        ):
+            # In case of vectorization, we keep the parent node name as, for
+            # instance, rate[status].zone1 is best described as the value of
+            # "rate".
+            name = self.parameter_node_at_instant._name
+
+        else:
+            name = f"{self.parameter_node_at_instant._name}.{key}"
+
+        if isinstance(child, (numpy.ndarray, *parameters.ALLOWED_PARAM_TYPES)):
+            self.tracer.record_parameter_access(name, period, child)
+
+        return child