Add tests for CSV format

Author: zooba

src/manage/list_command.py

@@ -105,33 +105,41 @@ def format_table(cmd, installs):
                      " for alternative ways to display this information.!W!")
 
 
-CSV_EXCLUDE = {
+CSV_EXCLUDE = frozenset([
     "schema", "unmanaged",
     # Complex columns of limited value
     "install-for", "shortcuts", "__original-shortcuts",
     "executable", "executable_args",
-}
+])
 
-CSV_EXPAND = ["run-for", "alias"]
+CSV_EXPAND = frozenset(["run-for", "alias"])
 
-def _csv_filter_and_expand(installs):
+def _csv_filter_and_expand(installs, *, exclude=CSV_EXCLUDE, expand=CSV_EXPAND):
     for i in installs:
         filtered = {}
-        to_expand = {k: [] for k in CSV_EXPAND}
+        to_expand = {k: [] for k in expand}
         for k, v in i.items():
-            if k in CSV_EXCLUDE:
+            if k in exclude:
                 continue
-            elif k in to_expand:
+            elif k in to_expand and isinstance(v, (list, tuple)):
                 for vv in v:
-                    expanded = {f"{k}.{k2}": vvv for k2, vvv in vv.items()}
+                    try:
+                        items = vv.items
+                    except AttributeError:
+                        expanded = {f"{k}": vv}
+                    else:
+                        expanded = {f"{k}.{k2}": vvv for k2, vvv in items()}
                     to_expand[k].append(expanded)
             else:
                 filtered[k] = v
 
-        yield filtered
-        for k in CSV_EXPAND:
+        any_yielded = False
+        for k in expand:
             for expanded in to_expand[k]:
                 yield filtered | expanded
+                any_yielded = True
+        if not any_yielded:
+            yield filtered
 
 
 def format_csv(cmd, installs):
@@ -140,7 +148,13 @@ def format_csv(cmd, installs):
     if not installs:
         return
     columns = list(dict.fromkeys(col for i in installs for col in i))
-    writer = csv.DictWriter(sys.stdout, columns)
+
+    class LoggingIOWrapper:
+        @staticmethod
+        def write(s):
+            LOGGER.print_raw(s, end="")
+
+    writer = csv.DictWriter(LoggingIOWrapper, columns)
     writer.writeheader()
     writer.writerows(installs)
 

tests/test_list.py

@@ -163,3 +163,63 @@ def test_format_table_empty(assert_log):
         (r"!B!Tag\s+Name\s+Managed By\s+Version\s+Alias\s*!W!", ()),
         (r".+No runtimes.+", ()),
     )
+
+
+def test_format_csv(assert_log):
+    list_command.format_csv(None, FAKE_INSTALLS)
+    # CSV format only contains columns that are present, so this doesn't look
+    # as complete as for normal installs, but it's fine for the test.
+    assert_log(
+        "company,tag,sort-version,default",
+        "Company2,1.0,1.0,",
+        "Company1,2.0,2.0,",
+        "Company1,1.0,1.0,True",
+    )
+
+
+def test_format_csv_complex(assert_log):
+    data = [
+        {
+            **d,
+            "alias": [dict(name=f"n{i}.{j}", target=f"t{i}.{j}") for j in range(i + 1)]
+        }
+        for i, d in enumerate(FAKE_INSTALLS)
+    ]
+    list_command.format_csv(None, data)
+    assert_log(
+        "company,tag,sort-version,alias.name,alias.target.default",
+        "Company2,1.0,1.0,n0.0,t0.0,",
+        "Company1,2.0,2.0,n1.0,t1.0,",
+        "Company1,2.0,2.0,n1.1,t1.1,",
+        "Company1,1.0,1.0,n2.0,t2.0,True",
+        "Company1,1.0,1.0,n2.1,t2.1,True",
+        "Company1,1.0,1.0,n2.2,t2.2,True",
+    )
+
+
+def test_format_csv_empty(assert_log):
+    list_command.format_csv(None, [])
+    assert_log(assert_log.end_of_log())
+
+
+def test_csv_exclude():
+    result = list(list_command._csv_filter_and_expand([
+        dict(a=1, b=2),
+        dict(a=3, c=4),
+        dict(a=5, b=6, c=7),
+    ], exclude={"b"}))
+    assert result == [dict(a=1), dict(a=3, c=4), dict(a=5, c=7)]
+
+
+def test_csv_expand():
+    result = list(list_command._csv_filter_and_expand([
+        dict(a=[1, 2], b=[3, 4]),
+        dict(a=[5], b=[6]),
+        dict(a=7, b=8),
+    ], expand={"a"}))
+    assert result == [
+        dict(a=1, b=[3, 4]),
+        dict(a=2, b=[3, 4]),
+        dict(a=5, b=[6]),
+        dict(a=7, b=8),
+    ]