Merge 6529d12 into b9ead83

openfisca_core/scripts/openfisca_command.py

@@ -41,6 +41,7 @@ def build_test_parser(parser):
         parser.add_argument('-v', '--verbose', action = 'store_true', default = False, help = "increase output verbosity")
         parser.add_argument('-o', '--only-variables', nargs = '*', default = None, help = "variables to test. If specified, only test the given variables.")
         parser.add_argument('-i', '--ignore-variables', nargs = '*', default = None, help = "variables to ignore. If specified, do not test the given variables.")
+        parser.add_argument('-d', '--depth_max', type = int, default = None, help = "Maximal depth of the log. If not specified, display all the computation")
 
         return parser
 

openfisca_core/scripts/run_test.py

@@ -22,6 +22,7 @@ def main(parser):
         'name_filter': args.name_filter,
         'only_variables': args.only_variables,
         'ignore_variables': args.ignore_variables,
+        'depth_max': args.depth_max,
         }
 
     paths = [os.path.abspath(path) for path in args.path]

openfisca_core/tools/test_runner.py

@@ -146,6 +146,7 @@ def runtest(self):
         verbose = self.options.get('verbose')
         performance_graph = self.options.get('performance_graph')
         performance_tables = self.options.get('performance_tables')
+        depth_max = self.options.get('depth_max')
 
         try:
             builder.set_default_period(period)
@@ -165,15 +166,15 @@ def runtest(self):
         finally:
             tracer = self.simulation.tracer
             if verbose:
-                self.print_computation_log(tracer)
+                self.print_computation_log(tracer, depth_max)
             if performance_graph:
                 self.generate_performance_graph(tracer)
             if performance_tables:
                 self.generate_performance_tables(tracer)
 
-    def print_computation_log(self, tracer):
+    def print_computation_log(self, tracer, depth_max = None):
         print("Computation log:")  # noqa T001
-        tracer.print_computation_log()
+        tracer.print_computation_log(depth_max = depth_max)
 
     def generate_performance_graph(self, tracer):
         tracer.generate_performance_graph('.')

openfisca_core/tracers/computation_log.py

@@ -35,6 +35,7 @@ def _get_node_log(
             node: tracers.TraceNode,
             depth: int,
             aggregate: bool,
+            depth_max: int = None,
             ) -> List[str]:
 
         def print_line(depth: int, node: tracers.TraceNode) -> str:
@@ -60,12 +61,16 @@ def print_line(depth: int, node: tracers.TraceNode) -> str:
 
             return f"{indent}{node.name}<{node.period}> >> {formatted_value}"
 
+        if depth_max is not None and depth > depth_max:
+            return
+
         node_log = [print_line(depth, node)]
 
         children_logs = [
-            self._get_node_log(child, depth + 1, aggregate)
+            self._get_node_log(child, depth + 1, aggregate, depth_max = depth_max)
             for child
             in node.children
+            if self._get_node_log(child, depth + 1, aggregate, depth_max = depth_max) is not None
             ]
 
         return node_log + self._flatten(children_logs)
@@ -76,18 +81,18 @@ def _flatten(
             ) -> List[str]:
         return [item for _list in list_of_lists for item in _list]
 
-    def lines(self, aggregate: bool = False) -> List[str]:
+    def lines(self, aggregate: bool = False, depth_max:int = None) -> List[str]:
         depth = 1
 
         lines_by_tree = [
-            self._get_node_log(node, depth, aggregate)
+            self._get_node_log(node, depth, aggregate, depth_max = depth_max)
             for node
             in self._full_tracer.trees
             ]
 
         return self._flatten(lines_by_tree)
 
-    def print_log(self, aggregate = False) -> None:
+    def print_log(self, aggregate = False, depth_max = None) -> None:
         """
         Print the computation log of a simulation.
 
@@ -99,5 +104,5 @@ def print_log(self, aggregate = False) -> None:
 
         This mode is more suited for simulations on a large population.
         """
-        for line in self.lines(aggregate):
+        for line in self.lines(aggregate, depth_max):
             print(line)  # noqa T001

openfisca_core/tracers/full_tracer.py

@@ -121,8 +121,8 @@ def flat_trace(self) -> tracers.FlatTrace:
     def _get_time_in_sec(self) -> float:
         return time.time_ns() / (10**9)
 
-    def print_computation_log(self, aggregate = False):
-        self.computation_log.print_log(aggregate)
+    def print_computation_log(self, depth_max = None, aggregate = False):
+        self.computation_log.print_log(aggregate, depth_max = depth_max)
 
     def generate_performance_graph(self, dir_path: str) -> None:
         self.performance_log.generate_graph(dir_path)