Improve CSV output (#44)

Co-authored-by: Steve Dower <[email protected]>

Author: AA-Turner

src/manage/list_command.py

@@ -105,34 +105,56 @@ 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:
-        i = {k: v for k, v in i.items() if k not in CSV_EXCLUDE}
-        to_expand = {k: i.pop(k, ()) for k in CSV_EXPAND}
-        yield i
-        for k2, vlist in to_expand.items():
-            for vv in vlist:
-                yield {f"{k2}.{k}": v for k, v in vv.items()}
+        filtered = {}
+        to_expand = {k: [] for k in expand}
+        for k, v in i.items():
+            if k in exclude:
+                continue
+            elif k in to_expand and isinstance(v, (list, tuple)):
+                for vv in v:
+                    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
+
+        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):
     import csv
     installs = list(_csv_filter_and_expand(installs))
     if not installs:
         return
-    s = set()
-    columns = [c for i in installs for c in i
-               if c not in s and (s.add(c) or True)]
-    writer = csv.DictWriter(sys.stdout, columns)
+    columns = list(dict.fromkeys(col for i in installs for col in i))
+
+    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),
+    ]