From baafed4a041886f3980de37a63294c57017cc5e7 Mon Sep 17 00:00:00 2001
From: Ramesh Padmanabhaiah <22363102+codeforester@users.noreply.github.com>
Date: Fri, 7 Aug 2026 11:40:26 -0700
Subject: [PATCH] Preserve heterogeneous document columns

---
 lib/python/base_cli/output.py | 12 +++++++---
 tests/test_output.py          | 44 ++++++++++++++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/lib/python/base_cli/output.py b/lib/python/base_cli/output.py
index 7dcee4b..3661c9e 100644
--- a/lib/python/base_cli/output.py
+++ b/lib/python/base_cli/output.py
@@ -171,9 +171,15 @@ def render_document(
 
 
 def _document_columns(records: Sequence[Mapping[str, Any]]) -> list[tuple[str, str]]:
-    if not records:
-        return []
-    return [(str(key).upper(), str(key)) for key in records[0]]
+    column_keys: list[str] = []
+    seen: set[str] = set()
+    for record in records:
+        for key in record:
+            normalized_key = str(key)
+            if normalized_key not in seen:
+                seen.add(normalized_key)
+                column_keys.append(normalized_key)
+    return [(key.upper(), key) for key in column_keys]
 
 
 def _cell_value(value: Any) -> str:
diff --git a/tests/test_output.py b/tests/test_output.py
index e3024cd..688bee7 100644
--- a/tests/test_output.py
+++ b/tests/test_output.py
@@ -4,7 +4,7 @@
 import json
 import unittest
 
-from base_cli.output import OutputFormatError, render_records, resolve_output_format
+from base_cli.output import OutputFormatError, render_document, render_records, resolve_output_format
 
 
 class _Stream(io.StringIO):
@@ -240,3 +240,45 @@ def test_empty_tty_result_keeps_footer(self) -> None:
         render_records((), requested_format="text", columns=COLUMNS, stream=stream, footer="0 projects.")
 
         self.assertEqual(stream.getvalue(), "0 projects.\n")
+
+    def test_render_document_unions_heterogeneous_record_columns(self) -> None:
+        stream = _Stream(terminal=False)
+        document = {
+            "rows": [
+                {"name": "one", "path": "/tmp/one"},
+                {"name": "two", "extra": "only later"},
+            ]
+        }
+
+        render_document(
+            document,
+            requested_format="csv",
+            records_key="rows",
+            stream=stream,
+        )
+
+        self.assertEqual(
+            stream.getvalue(),
+            "one,/tmp/one,\ntwo,,only later\n",
+        )
+
+    def test_render_document_tsv_includes_later_record_columns(self) -> None:
+        stream = _Stream(terminal=False)
+        document = {
+            "rows": [
+                {"name": "one"},
+                {"name": "two", "status": "ready"},
+            ]
+        }
+
+        render_document(
+            document,
+            requested_format="tsv",
+            records_key="rows",
+            stream=stream,
+        )
+
+        output = stream.getvalue()
+        self.assertIn("one\t\n", output)
+        self.assertIn("two\tready\n", output)
+        self.assertIn("ready", output)