repl mode

Author: brunovcosta

json_sql/cli.py

@@ -4,6 +4,7 @@
 from argparse import ArgumentParser
 from csv import DictWriter
 from sys import stdout
+from typing import List, Literal
 from json import dumps
 
 
@@ -13,9 +14,23 @@ def query(code: str, workdir: Path, ctx: dict):
     return result
 
 
+def print_result(result: List[dict], format: Literal["json", "csv"]):
+    if format == "json":
+        print(dumps(result, indent=2))
+    elif format == "csv":
+        if len(result) == 0:
+            print("No results to write to CSV.")
+        else:
+            keys = result[0].keys()
+            writer = DictWriter(stdout, fieldnames=keys)
+            writer.writeheader()
+            for row in result:
+                writer.writerow(row)
+
+
 if __name__ == "__main__":
     parser = ArgumentParser(description="Run SQL queries on JSON files.")
-    parser.add_argument("code", type=str, help="SQL query to execute")
+    parser.add_argument("--code", type=str, help="SQL query to execute", default=None)
     parser.add_argument(
         "--workdir",
         type=Path,
@@ -37,16 +52,25 @@ def query(code: str, workdir: Path, ctx: dict):
     )
     args = parser.parse_args()
 
-    result = query(code=args.code, workdir=args.workdir, ctx={})
-
-    if args.format == "json":
-        print(dumps(result, indent=2))
-    elif args.format == "csv":
-        if len(result) == 0:
-            print("No results to write to CSV.")
-        else:
-            keys = result[0].keys()
-            writer = DictWriter(stdout, fieldnames=keys)
-            writer.writeheader()
-            for row in result:
-                writer.writerow(row)
+    if args.code:
+        result = query(code=args.code, workdir=args.workdir, ctx={})
+        print_result(result=result, format=args.format)
+    else:
+        print("JSON SQL CLI")
+        print("Type 'exit' to quit.")
+        while True:
+            try:
+                code = input("> ")
+                if code.lower() == "exit":
+                    break
+                result = query(code=code, workdir=args.workdir, ctx={})
+                print_result(result=result, format=args.format)
+            except KeyboardInterrupt:
+                break
+            except EOFError:
+                break
+            except SystemExit:
+                break
+            except Exception as e:
+                print(f"Error: {e}")
+                continue