Add --handle-json to codegen

edgedb/codegen/cli.py

@@ -60,6 +60,12 @@
         help="Add a mixin to generated dataclasses "
         "to skip Pydantic validation (default is to add the mixin).",
     )
+    parser.add_argument(
+        "--handle-json",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Choose to handle JSON in query arguments and results."
+    )
 else:
     parser.add_argument(
         "--skip-pydantic-validation",
@@ -74,6 +80,19 @@
         help="Add a mixin to generated dataclasses "
              "to skip Pydantic validation (default is to add the mixin).",
     )
+    parser.add_argument(
+        "--handle-json",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--no-handle-json",
+        dest="handle_json",
+        action="store_false",
+        default=True,
+        help="Choose to handle JSON in query arguments and results. "
+             "(default: False)"
+    )
 
 
 def main():

edgedb/codegen/generator.py

@@ -123,6 +123,7 @@ def __init__(self, args: argparse.Namespace):
         self._default_module = "default"
         self._targets = args.target
         self._skip_pydantic_validation = args.skip_pydantic_validation
+        self._handle_json = args.handle_json
         self._async = False
         try:
             self._project_dir = pathlib.Path(
@@ -339,6 +340,16 @@ def _generate(
             method = "query_single"
             rt = "return "
 
+        if self._handle_json:
+            print(f"{INDENT}client = client.with_codec_context(", file=buf)
+            print(
+                f"{INDENT}{INDENT}"
+                f"edgedb.get_default_codec_context(handle_json=True),",
+                file=buf,
+            )
+            print(f"{INDENT}{INDENT}only_replace_default=True,", file=buf)
+            print(f"{INDENT})", file=buf)
+
         if self._async:
             print(f"{INDENT}{rt}await client.{method}(", file=buf)
         else:
@@ -369,7 +380,17 @@ def _generate_code(
         if type_.desc_id in self._cache:
             return self._cache[type_.desc_id]
 
-        if isinstance(type_, describe.BaseScalarType):
+        if (
+            isinstance(type_, describe.BaseScalarType)
+            and type_.name == "std::json"
+        ):
+            if self._handle_json:
+                self._imports.add("typing")
+                rv = "typing.Any"
+            else:
+                rv = "str"
+
+        elif isinstance(type_, describe.BaseScalarType):
             if type_.name in TYPE_IMPORTS:
                 self._imports.add(TYPE_IMPORTS[type_.name])
             rv = TYPE_MAPPING[type_.name]

edgedb/options.py

@@ -370,6 +370,8 @@ def without_globals(self, *global_names):
     def with_codec_context(
         self, codec_context: pgproto.CodecContext, only_replace_default=False
     ):
+        if self._options.codec_context is codec_context:
+            return self
         if only_replace_default:
             default = protocol.get_default_codec_context()
             if self._options.codec_context is not default:

tests/codegen/test-project1/generated_async_edgeql.py.assert

@@ -1,5 +1,5 @@
 # AUTOGENERATED FROM 'select_scalar.edgeql' WITH:
-#     $ edgedb-py --target async --file --no-skip-pydantic-validation
+#     $ edgedb-py --target async --file --no-skip-pydantic-validation --handle-json
 
 
 from __future__ import annotations
@@ -9,6 +9,10 @@ import edgedb
 async def select_scalar(
     client: edgedb.AsyncIOClient,
 ) -> int:
+    client = client.with_codec_context(
+        edgedb.get_default_codec_context(handle_json=True),
+        only_replace_default=True,
+    )
     return await client.query_single(
         """\
         select 1;\

tests/codegen/test-project2/generated_async_edgeql.py.assert

@@ -8,7 +8,7 @@
 #     'scalar/select_scalar.edgeql'
 #     'scalar/select_scalars.edgeql'
 # WITH:
-#     $ edgedb-py --target async --file --no-skip-pydantic-validation
+#     $ edgedb-py --target async --file --no-skip-pydantic-validation --handle-json
 
 
 from __future__ import annotations
@@ -137,6 +137,10 @@ class SelectObjectResultParamsItem:
 async def link_prop(
     client: edgedb.AsyncIOClient,
 ) -> list[LinkPropResult]:
+    client = client.with_codec_context(
+        edgedb.get_default_codec_context(handle_json=True),
+        only_replace_default=True,
+    )
     return await client.query(
         """\
         create type Person {
@@ -215,6 +219,10 @@ async def my_query(
     aw: edgedb.Range[datetime.date],
     ax: edgedb.Range[datetime.date],
 ) -> MyQueryResult:
+    client = client.with_codec_context(
+        edgedb.get_default_codec_context(handle_json=True),
+        only_replace_default=True,
+    )
     return await client.query_single(
         """\
         create scalar type MyScalar extending int64;
@@ -335,6 +343,10 @@ async def query_one(
     *,
     arg_name_with_underscores: int,
 ) -> int:
+    client = client.with_codec_context(
+        edgedb.get_default_codec_context(handle_json=True),
+        only_replace_default=True,
+    )
     return await client.query_single(
         """\
         select <int64>$arg_name_with_underscores\
@@ -349,6 +361,10 @@ async def select_args(
     arg_str: str,
     arg_datetime: datetime.datetime,
 ) -> SelectArgsResult:
+    client = client.with_codec_context(
+        edgedb.get_default_codec_context(handle_json=True),
+        only_replace_default=True,
+    )
     return await client.query_single(
         """\
         select {
@@ -364,6 +380,10 @@ async def select_args(
 async def select_object(
     client: edgedb.AsyncIOClient,
 ) -> SelectObjectResult | None:
+    client = client.with_codec_context(
+        edgedb.get_default_codec_context(handle_json=True),
+        only_replace_default=True,
+    )
     return await client.query_single(
         """\
         select schema::Function {
@@ -382,6 +402,10 @@ async def select_object(
 async def select_objects(
     client: edgedb.AsyncIOClient,
 ) -> list[SelectObjectResult]:
+    client = client.with_codec_context(
+        edgedb.get_default_codec_context(handle_json=True),
+        only_replace_default=True,
+    )
     return await client.query(
         """\
         select schema::Function {
@@ -399,6 +423,10 @@ async def select_objects(
 async def select_scalar(
     client: edgedb.AsyncIOClient,
 ) -> int:
+    client = client.with_codec_context(
+        edgedb.get_default_codec_context(handle_json=True),
+        only_replace_default=True,
+    )
     return await client.query_single(
         """\
         select 1;\
@@ -409,6 +437,10 @@ async def select_scalar(
 async def select_scalars(
     client: edgedb.AsyncIOClient,
 ) -> list[edgedb.ConfigMemory]:
+    client = client.with_codec_context(
+        edgedb.get_default_codec_context(handle_json=True),
+        only_replace_default=True,
+    )
     return await client.query(
         """\
         select <cfg::memory>{1, 2, 3};\

tests/test_codegen.py

@@ -82,6 +82,7 @@ async def run(*args, extra_env=None):
             "async",
             "--file",
             "--no-skip-pydantic-validation",
+            "--handle-json",
             extra_env={"EDGEDB_PYTHON_CODEGEN_PY_VER": "3.10.3"},
         )