diff --git a/src/manage/list_command.py b/src/manage/list_command.py index fde28bb..a13c787 100644 --- a/src/manage/list_command.py +++ b/src/manage/list_command.py @@ -105,23 +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: - 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): @@ -129,10 +147,14 @@ def format_csv(cmd, installs): 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) diff --git a/tests/test_list.py b/tests/test_list.py index 12003af..f2b5ace 100644 --- a/tests/test_list.py +++ b/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), + ]