remove print statements, expand test coverage

NYCPlanning · Oct 8, 2024 · 67266f6 · 67266f6
1 parent 2d6e745
commit 67266f6
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 51 deletions.
diff --git a/dcpy/test/utils/test_collections.py b/dcpy/test/utils/test_collections.py
@@ -14,8 +14,8 @@ def test_deep_merge():
 class Inner(BaseModel, arbitrary_types_allowed=True):
     g: str
     i: list
-    l: dict[str, str]  # noqa: E741
-    q: pd.DataFrame
+    m: dict[str, str]
+    r: pd.DataFrame
 
 
 class Outer(BaseModel, arbitrary_types_allowed=True):
@@ -27,32 +27,36 @@ class Outer(BaseModel, arbitrary_types_allowed=True):
 
 class TestIndentedPrint:
     df = pd.DataFrame({"x": [1, 2], "y": [3, 4]})
-    inner = Inner(g="h", i=["j", "k"], l={"m": "n", "o": "p"}, q=df)
+    inner = Inner(g="h", i=[["j"], {"k", "l"}], m={"n": "o", "p": "q"}, r=df)
     outer = Outer(a={"b", "c"}, d=df, e={}, f=inner)
 
     def test(self):
-        assert collections.flatten_and_indent(self.outer.model_dump()) == [
-            "a",
-            "    b",
-            "    c",
-            "d",
-            "       x  y",
-            "    0  1  3",
-            "    1  2  4",
-            "e: None",
-            "f",
-            "    g: h",
-            "    i",
-            "        j",
-            "        k",
-            "    l",
-            "        m: n",
-            "        o: p",
-            "    q",
-            "           x  y",
-            "        0  1  3",
-            "        1  2  4",
-        ]
+        assert (
+            collections.flatten_and_indent(self.outer.model_dump())
+            == [
+                "a",
+                "    b",
+                "    c",
+                "d",
+                "       x  y",
+                "    0  1  3",
+                "    1  2  4",
+                "e: None",
+                "f",
+                "    g: h",
+                "    i",
+                "            j",  # list of lists are a little awkward. Not really a common use case though
+                "            k",
+                "            l",
+                "    m",
+                "        n: o",
+                "        p: q",
+                "    r",
+                "           x  y",
+                "        0  1  3",
+                "        1  2  4",
+            ]
+        )
 
     def test_custom_indent_and_offset(self):
         assert collections.flatten_and_indent(
@@ -69,12 +73,13 @@ def test_custom_indent_and_offset(self):
             " f",
             "   g: h",
             "   i",
-            "     j",
-            "     k",
-            "   l",
-            "     m: n",
-            "     o: p",
-            "   q",
+            "       j",
+            "       k",
+            "       l",
+            "   m",
+            "     n: o",
+            "     p: q",
+            "   r",
             "        x  y",
             "     0  1  3",
             "     1  2  4",
@@ -99,12 +104,13 @@ def test_line_breaks(self):
             "f",
             "    g: h",
             "    i",
-            "        j",
-            "        k",
-            "    l",
-            "        m: n",
-            "        o: p",
-            "    q",
+            "            j",
+            "            k",
+            "            l",
+            "    m",
+            "        n: o",
+            "        p: q",
+            "    r",
             "           x  y",
             "        0  1  3",
             "        1  2  4",
@@ -125,12 +131,13 @@ def test_pretty_print(self):
             "F",
             "    G: h",
             "    I",
-            "        j",
-            "        k",
-            "    L",
-            "        M: n",
-            "        O: p",
-            "    Q",
+            "            j",
+            "            k",
+            "            l",
+            "    M",
+            "        N: o",
+            "        P: q",
+            "    R",
             "           x  y",
             "        0  1  3",
             "        1  2  4",
@@ -148,12 +155,13 @@ def test_to_report(self):
         f
             g: h
             i
-                j
-                k
-            l
-                m: n
-                o: p
-            q
+                    j
+                    k
+                    l
+            m
+                n: o
+                p: q
+            r
                    x  y
                 0  1  3
                 1  2  4"""
diff --git a/dcpy/utils/collections.py b/dcpy/utils/collections.py
@@ -76,11 +76,9 @@ def flatten(li: list | dict, level: int) -> list:
         current_node = li[0]
 
         if isinstance(current_node, list) or isinstance(current_node, dict):
-            print("node is list or dict")
             first = flatten(current_node, level + 1)
 
         elif isinstance(current_node, set):
-            print("node is set")
             first = flatten(sorted(list(current_node)), level + 1)
 
         elif isinstance(current_node, DataFrame):