Add field selection to CSV and JSON outputs (#307)

jazzband · Aug 12, 2024 · 5fb898f · 5fb898f
1 parent e26318f
commit 5fb898f
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 6 deletions.
diff --git a/src/prettytable/prettytable.py b/src/prettytable/prettytable.py
@@ -2119,8 +2119,20 @@ def get_csv_string(self, **kwargs) -> str:
         csv_writer = csv.writer(csv_buffer, **csv_options)
 
         if options.get("header"):
-            csv_writer.writerow(self._field_names)
-        for row in self._get_rows(options):
+            if options["fields"]:
+                csv_writer.writerow(
+                    [f for f in self._field_names if f in options["fields"]]
+                )
+            else:
+                csv_writer.writerow(self._field_names)
+
+        rows = self._get_rows(options)
+        if options["fields"]:
+            rows = [
+                [d for f, d in zip(self._field_names, row) if f in options["fields"]]
+                for row in rows
+            ]
+        for row in rows:
             csv_writer.writerow(row)
 
         return csv_buffer.getvalue()
@@ -2144,12 +2156,26 @@ def get_json_string(self, **kwargs) -> str:
         json_options.update(
             {key: value for key, value in kwargs.items() if key not in options}
         )
-        objects = []
+        objects: list[list[str] | dict[str, Any]] = []
 
         if options.get("header"):
-            objects.append(self.field_names)
-        for row in self._get_rows(options):
-            objects.append(dict(zip(self._field_names, row)))
+            if options["fields"]:
+                objects.append([f for f in self._field_names if f in options["fields"]])
+            else:
+                objects.append(self.field_names)
+        rows = self._get_rows(options)
+        if options["fields"]:
+            for row in rows:
+                objects.append(
+                    {
+                        f: d
+                        for f, d in zip(self._field_names, row)
+                        if f in options["fields"]
+                    }
+                )
+        else:
+            for row in rows:
+                objects.append(dict(zip(self._field_names, row)))
 
         return json.dumps(objects, **json_options)
 

diff --git a/tests/test_prettytable.py b/tests/test_prettytable.py
@@ -1021,6 +1021,30 @@ def test_json_output(self) -> None:
         "Field 2": "value8",
         "Field 3": "value9"
     }
+]""".strip()
+        )
+        options = {"fields": ["Field 1", "Field 3"]}
+        result = t.get_json_string(**options)
+        assert (
+            result.strip()
+            == """
+[
+    [
+        "Field 1",
+        "Field 3"
+    ],
+    {
+        "Field 1": "value 1",
+        "Field 3": "value3"
+    },
+    {
+        "Field 1": "value 4",
+        "Field 3": "value6"
+    },
+    {
+        "Field 1": "value 7",
+        "Field 3": "value9"
+    }
 ]""".strip()
         )
 
@@ -1750,6 +1774,13 @@ def test_csv_output(self) -> None:
             "4,value 4,value5,value6\r\n"
             "7,value 7,value8,value9\r\n"
         )
+        options = {"fields": ["Field 1", "Field 3"]}
+        assert t.get_csv_string(**options) == (
+            "Field 1,Field 3\r\n"
+            "value 1,value3\r\n"
+            "value 4,value6\r\n"
+            "value 7,value9\r\n"
+        )
 
 
 class TestLatexOutput: