extractor works

extract.py

@@ -3,7 +3,6 @@
 import ast
 import argparse
 import json
-import re
 from pathlib import Path
 
 DEFAULT_REPO = "/home/uha4/tmp/tryton"
@@ -29,25 +28,9 @@ class RegisterExtractor(ast.NodeVisitor):
 
     def extract_register(self, cls, func):
 
-        phase = "before_super"
-
         operations = []
 
-        for stmt in func.body:
-
-            # Wurde super().__register__() aufgerufen?
-            if self.is_super_register(stmt):
-                phase = "after_super"
-                continue
-
-            for call in ast.walk(stmt):
-                if isinstance(call, ast.Call):
-
-                    op = self.extract_call(call)
-
-                    if op:
-                        op["phase"] = phase
-                        operations.append(op)
+        self.walk_statements(func.body, operations, "before_super", None)
 
         self.results.append({
             "file": self.filename,
@@ -57,6 +40,41 @@ class RegisterExtractor(ast.NodeVisitor):
             "operations": operations,
         })
 
+
+    def walk_statements(self, statements, operations, phase, condition):
+
+        for stmt in statements:
+
+            if self.is_super_register(stmt):
+                phase = "after_super"
+                continue
+
+            if isinstance(stmt, ast.If):
+                cond = ast.unparse(stmt.test)
+                self.walk_statements(stmt.body, operations, phase, cond)
+                self.walk_statements(stmt.orelse, operations, phase, condition)
+                continue
+
+            value = getattr(stmt, "value", None)
+            if value is not None:
+                self.walk_expression(value, operations, phase, condition)
+
+    def walk_expression(self, node, operations, phase, condition):
+
+        if node is None:
+            return
+
+        if isinstance(node, ast.Call):
+            op = self.extract_call(node)
+            if op:
+                op["phase"] = phase
+                if condition:
+                    op["condition"] = condition
+                operations.append(op)
+
+        for child in ast.iter_child_nodes(node):
+            self.walk_expression(child, operations, phase, condition)
+
     def is_super_register(self, stmt):
 
         if not isinstance(stmt, ast.Expr):