refactor: refactor Forth generation

src/uproot/_awkward_forth.py

@@ -6,8 +6,6 @@
 
 import numpy as np
 
-import uproot.containers
-
 symbol_dict = {
     np.dtype(">f4"): "f",
     np.dtype(">f8"): "d",
@@ -22,14 +20,6 @@
 }
 
 
-def check_depth(node):
-    """
-    This method checks the depth of the provided container
-    """
-    if isinstance(node, uproot.containers.AsVector):
-        return 1
-
-
 class ForthGenerator:
     """
     This class is passed through the Forth code generation, collecting Forth snippets and concatenating them at the end.
@@ -43,7 +33,7 @@ def __init__(self, aform=None, count_obj=0, var_set=False):
         self.top_form = None
         self.prev_form = None
         self.awkward_model = {"name": "TOP", "content": {}}
-        self._prev_node = self.awkward_model
+        self.top_node = self.awkward_model
         self.ref_list = []
         self.forth_code = {}
         self.forth_keys = {}
@@ -61,9 +51,9 @@ def traverse_aform(self):
     def replace_form_and_model(self, form, model):
         temp_node = self.awkward_model
         temp_prev_form = self.prev_form
-        temp_node_top = self._prev_node
+        temp_node_top = self.top_node
         self.awkward_model = model
-        self._prev_node = self.awkward_model
+        self.top_node = self.awkward_model
         temp_form = self.aform
         temp_form_top = self.top_form
         self.top_form = None
@@ -106,83 +96,41 @@ def should_add_form(self):
     def get_temp_form_top(self):
         return self.top_dummy
 
-    def set_dummy_none(self, temp_top, temp_form, temp_flag):
-        self.top_dummy = temp_top
-        self.dummy_aform = temp_form
-        self.dummy_form = temp_flag
-
     def add_form(self, aform, conlen=0, traverse=1):
-        if self.dummy_form:
-            if self.dummy_aform is None:
-                self.dummy_aform = aform
-                self.top_dummy = aform
-            else:
-                if "content" in self.dummy_aform.keys():
-                    if self.dummy_aform["content"] == "NULL":
-                        self.dummy_aform["content"] = aform
-                        self.dummy_aform = self.dummy_aform["content"]
-                    else:
-                        raise ValueError
-                elif "contents" in self.dummy_aform.keys():
-                    if (
-                        len(self.dummy_aform["content"])
-                        < self.dummy_aform["parameters"]["lencon"]
-                    ):
-                        self.dummy_aform["contents"].append(aform)
-                    else:
-                        raise ValueError
+        if self.aform is None:
+            self.aform = aform
+            self.top_form = self.aform
+            if traverse == 2:
+                self.aform = self.aform["content"]
         else:
-            if self.aform is None:
-                self.aform = aform
-                self.top_form = self.aform
-                if traverse == 2:
-                    self.aform = self.aform["content"]
-            else:
-                if "content" in self.aform.keys():
-                    if self.aform["content"] == "NULL":
-                        self.aform["content"] = aform
-                        self.prev_form = self.aform
-                        if traverse == 2:
-                            self.aform = self.aform["content"]["content"]
-                            self.prev_form = self.prev_form["content"]
-                        else:
-                            self.aform = self.aform["content"]
+            if "content" in self.aform.keys():
+                if self.aform["content"] == "NULL":
+                    self.aform["content"] = aform
+                    self.prev_form = self.aform
+                    if traverse == 2:
+                        self.aform = self.aform["content"]["content"]
+                        self.prev_form = self.prev_form["content"]
                     else:
-                        raise ValueError
-                elif "contents" in self.aform.keys():
-                    if self.aform["class"] == "RecordArray":
-                        if self.prev_form is not None:
-                            self.prev_form["content"] = aform
-                            self.aform = aform
-                        else:
-                            self.top_form = aform
-                            self.aform = aform
-                        # for elem in aform['contents'].keys():
-                        #     if self.depth(self.aform['contents'][elem])< self.depth(aform['contents'][elem]):
-                        #         aform['contents'][elem] = self.replace_keys(self.aform['contents'][elem],aform['contents'][elem])
-                        #         self.aform['contents'][elem] = aform['contents'][elem]
-                    elif len(self.aform["contents"]) == conlen:
-                        pass
+                        self.aform = self.aform["content"]
+                else:
+                    raise AssertionError("Form representation is corrupted.")
+            elif "contents" in self.aform.keys():
+                if self.aform["class"] == "RecordArray":
+                    if self.prev_form is not None:
+                        self.prev_form["content"] = aform
+                        self.aform = aform
                     else:
-                        raise ValueError
-
-    def depth(self, form):
-        count = 0
-        temp = form
-        while "content" in temp.keys():
-            count += 1
-            if isinstance(temp["content"], str):
-                break
-            temp = temp["content"]
-        return count
-
-    def replace_keys(self, old, new):
-        temp = new
-        while old != "NULL":
-            new["form_key"] = old["form_key"]
-            new = new["content"]
-            old = old["content"]
-        return temp
+                        self.top_form = aform
+                        self.aform = aform
+                elif len(self.aform["contents"]) == conlen:
+                    pass
+                else:
+                    return False
+
+    def set_dummy_none(self, temp_top, temp_form, temp_flag):
+        self.top_dummy = temp_top
+        self.dummy_aform = temp_form
+        self.dummy_form = temp_flag
 
     def get_keys(self, num_keys):
         if num_keys == 1:
@@ -196,13 +144,12 @@ def get_keys(self, num_keys):
                 self.count_obj += 1
             return out
         else:
-            raise ValueError("Number of keys cannot be less than 1")
+            raise AssertionError("Number of keys cannot be less than 1")
 
     def add_form_key(self, key):
         self.form_keys.append(key)
 
     def go_to(self, aform):
-        # aform["content"] = self.awkward_model
         self.awkward_model = aform
 
     def become(self, aform):
@@ -212,7 +159,7 @@ def check_model(self):
         return bool(self.awkward_model)
 
     def get_current_node(self):
-        self.ref_list.append(self._prev_node)
+        self.ref_list.append(self.top_node)
         return len(self.ref_list) - 1
 
     def get_ref(self, index):
@@ -230,9 +177,7 @@ def add_node_whole(self, new_node, ref_latest):
             self.awkward_model["contents"].append(new_node)
         self.awkward_model = ref_latest
 
-    def add_node(
-        self, name, precode, postcode, initcode, headercode, dtype, num_child, content
-    ):
+    def add_node(self, name, code_attrs, dtype, num_child, content):
         if isinstance(self.awkward_model, dict):
             if (
                 not bool(self.awkward_model["content"])
@@ -241,10 +186,10 @@ def add_node(
                 temp_obj = {
                     "name": name,
                     "type": dtype,
-                    "pre_code": precode,
-                    "post_code": postcode,
-                    "init_code": initcode,
-                    "header_code": headercode,
+                    "pre_code": code_attrs["precode"],
+                    "post_code": code_attrs["postcode"],
+                    "init_code": code_attrs["initcode"],
+                    "header_code": code_attrs["headercode"],
                     "num_child": num_child,
                     "content": content,
                 }
@@ -257,10 +202,10 @@ def add_node(
                         temp_obj = {
                             "name": name,
                             "type": dtype,
-                            "pre_code": precode,
-                            "post_code": postcode,
-                            "init_code": initcode,
-                            "header_code": headercode,
+                            "pre_code": code_attrs["precode"],
+                            "post_code": code_attrs["postcode"],
+                            "init_code": code_attrs["initcode"],
+                            "header_code": code_attrs["headercode"],
                             "num_child": num_child,
                             "content": content,
                         }
@@ -283,10 +228,10 @@ def add_node(
                 {
                     "name": name,
                     "type": dtype,
-                    "pre_code": precode,
-                    "post_code": postcode,
-                    "init_code": initcode,
-                    "header_code": headercode,
+                    "pre_code": code_attrs["precode"],
+                    "post_code": code_attrs["postcode"],
+                    "init_code": code_attrs["initcode"],
+                    "header_code": code_attrs["headercode"],
                     "num_child": num_child,
                     "content": content,
                 }
@@ -320,7 +265,7 @@ def add_forth_code(self, ref, forth_exec_pre, forth_exec_post):
 
     def add_to_final(self, code):
         if not isinstance(code, list):
-            raise TypeError
+            raise AssertionError("Something went wrong with Forth code generation.")
         self.final_code.extend(code)
         return
 
@@ -358,27 +303,23 @@ def get_gen_obj(self):
     def add_to_pre(self, code):
         self._pre_code.append(code)
 
-    def get_pre(self):
-        return self._pre_code
-
     def add_to_post(self, code):
         self._post_code.append(code)
 
-    def get_post(self):
-        return self._post_code
+    def get_attrs(self):
+        return {
+            "precode": self._pre_code,
+            "postcode": self._post_code,
+            "initcode": self._init,
+            "headercode": self._header,
+        }
 
     def add_to_header(self, code):
         self._header += code
 
-    def get_header(self):
-        return self._header
-
     def add_to_init(self, code):
         self._init += code
 
-    def get_init(self):
-        return self._init
-
 
 def convert_dtype(format):
     """Takes datatype codes from classses and returns the full datatype name.

src/uproot/containers.py

@@ -76,10 +76,7 @@ def _read_nested(
                 )
             temp = forth_obj.add_node(
                 f"node{key}",
-                helper_obj.get_pre(),
-                helper_obj.get_post(),
-                helper_obj.get_init(),
-                helper_obj.get_header(),
+                helper_obj.get_attrs(),
                 "i64",
                 1,
                 None,
@@ -546,10 +543,7 @@ def read(self, chunk, cursor, context, file, selffile, parent, header=True):
             helper_obj.add_to_init(f"0 node{offsets_num}-offsets <- stack\n")
             temp_form = forth_obj.add_node(
                 f"node{offsets_num}",
-                helper_obj.get_pre(),
-                helper_obj.get_post(),
-                helper_obj.get_init(),
-                helper_obj.get_header(),
+                helper_obj.get_attrs(),
                 "i64",
                 0,
                 None,
@@ -787,10 +781,7 @@ def read(self, chunk, cursor, context, file, selffile, parent, header=True):
                             forth_obj.add_form(json.loads(temp_aform))
                     temp = forth_obj.add_node(
                         f"node{offsets_num}",
-                        helper_obj.get_pre(),
-                        helper_obj.get_post(),
-                        helper_obj.get_init(),
-                        helper_obj.get_header(),
+                        helper_obj.get_attrs(),
                         "i64",
                         1,
                         {},
@@ -868,10 +859,7 @@ def read(self, chunk, cursor, context, file, selffile, parent, header=True):
                             forth_obj.add_form(json.loads(temp_aform))
                     temp = forth_obj.add_node(
                         f"node{offsets_num}",
-                        helper_obj.get_pre(),
-                        helper_obj.get_post(),
-                        helper_obj.get_init(),
-                        helper_obj.get_header(),
+                        helper_obj.get_attrs(),
                         "i64",
                         1,
                         {},
@@ -1197,10 +1185,7 @@ def read(self, chunk, cursor, context, file, selffile, parent, header=True):
                     helper_obj.add_to_post("loop\n")
                 temp = forth_obj.add_node(
                     f"node{key}",
-                    helper_obj.get_pre(),
-                    helper_obj.get_post(),
-                    helper_obj.get_init(),
-                    helper_obj.get_header(),
+                    helper_obj.get_attrs(),
                     "i64",
                     1,
                     {},
@@ -1363,10 +1348,7 @@ def read(self, chunk, cursor, context, file, selffile, parent, header=True):
                 helper_obj.add_to_post("loop\n")
             temp = forth_obj.add_node(
                 f"node{key}",
-                helper_obj.get_pre(),
-                helper_obj.get_post(),
-                helper_obj.get_init(),
-                helper_obj.get_header(),
+                helper_obj.get_attrs(),
                 "i64",
                 1,
                 {},
@@ -1572,7 +1554,7 @@ def read(self, chunk, cursor, context, file, selffile, parent, header=True):
             if helper_obj.is_forth():
                 temp = {"name": "TOP", "content": {}}
                 keys_form = forth_obj.top_form
-                keys_model = forth_obj._prev_node
+                keys_model = forth_obj.top_node
                 (
                     temp_node1,
                     temp_node_top1,
@@ -1604,11 +1586,11 @@ def read(self, chunk, cursor, context, file, selffile, parent, header=True):
             )
             if helper_obj.is_forth():
                 values_form = forth_obj.top_form
-                values_model = forth_obj._prev_node
+                values_model = forth_obj.top_node
                 if not isinstance(self._values, numpy.dtype):
                     values_model["content"]["post_code"].append("loop \n")
                 forth_obj.awkward_model = temp_node
-                forth_obj._prev_node = temp_node_top
+                forth_obj.top_node = temp_node_top
                 forth_obj.aform = temp_form
                 forth_obj.top_form = temp_form_top
                 forth_obj.prev_form = temp_prev_form
@@ -1627,10 +1609,7 @@ def read(self, chunk, cursor, context, file, selffile, parent, header=True):
                     forth_obj.add_form(aform)
                 temp = forth_obj.add_node(
                     "nodeMap",
-                    helper_obj.get_pre(),
-                    helper_obj.get_post(),
-                    helper_obj.get_init(),
-                    helper_obj.get_header(),
+                    helper_obj.get_attrs(),
                     "i64",
                     1,
                     [keys_model["content"], values_model["content"]],