fix: none-record queries

metatypedev · Jan 20, 2025 · 42782ca · 42782ca
1 parent 69bce82
commit 42782ca
Show file tree

Hide file tree

Showing 12 changed files with 442 additions and 33 deletions.
diff --git a/src/metagen/src/client_py/static/client.py b/src/metagen/src/client_py/static/client.py
@@ -196,6 +196,7 @@ def selection_to_nodes(
 #
 
 Out = typing.TypeVar("Out", covariant=True)
+PreparedOut = typing.TypeVar("PreparedOut", covariant=True)
 
 T = typing.TypeVar("T")
 
@@ -757,31 +758,80 @@ def mutation(
         result = self.fetch(doc, variables, opts, files)
         return result["value"] if isinstance(inp, MutationNode) else result
 
+    @typing.overload
+    def prepare_query(
+        self,
+        fun: typing.Callable[[PreparedArgs], QueryNode[Out]],
+        name: str = "",
+    ) -> "PreparedRequestUrlib[Out, Out]": ...
+
+    @typing.overload
     def prepare_query(
         self,
         fun: typing.Callable[[PreparedArgs], typing.Dict[str, QueryNode[Out]]],
         name: str = "",
-    ) -> "PreparedRequestUrlib[Out]":
+    ) -> "PreparedRequestUrlib[Out, typing.Dict[str, Out]]": ...
+
+    def prepare_query(
+        self,
+        fun: typing.Callable[
+            [PreparedArgs],
+            typing.Union[QueryNode[Out], typing.Dict[str, QueryNode[Out]]],
+        ],
+        name: str = "",
+    ) -> typing.Union[
+        "PreparedRequestUrlib[Out, Out]",
+        "PreparedRequestUrlib[Out, typing.Dict[str, Out]]",
+    ]:
         return PreparedRequestUrlib(self, fun, "query", name)
 
+    @typing.overload
+    def prepare_mutation(
+        self,
+        fun: typing.Callable[[PreparedArgs], MutationNode[Out]],
+        name: str = "",
+    ) -> "PreparedRequestUrlib[Out, Out]": ...
+
+    @typing.overload
     def prepare_mutation(
         self,
         fun: typing.Callable[[PreparedArgs], typing.Dict[str, MutationNode[Out]]],
         name: str = "",
-    ) -> "PreparedRequestUrlib[Out]":
+    ) -> "PreparedRequestUrlib[Out, typing.Dict[str, Out]]": ...
+
+    def prepare_mutation(
+        self,
+        fun: typing.Callable[
+            [PreparedArgs],
+            typing.Union[MutationNode[Out], typing.Dict[str, MutationNode[Out]]],
+        ],
+        name: str = "",
+    ) -> typing.Union[
+        "PreparedRequestUrlib[Out, Out]",
+        "PreparedRequestUrlib[Out, typing.Dict[str, Out]]",
+    ]:
         return PreparedRequestUrlib(self, fun, "mutation", name)
 
 
-class PreparedRequestBase(typing.Generic[Out]):
+class PreparedRequestBase(typing.Generic[Out, PreparedOut]):
     def __init__(
         self,
         transport: GraphQLTransportBase,
-        fun: typing.Callable[[PreparedArgs], typing.Mapping[str, SelectNode[Out]]],
+        fun: typing.Callable[
+            [PreparedArgs],
+            typing.Union[SelectNode[Out], typing.Mapping[str, SelectNode[Out]]],
+        ],
         ty: typing.Union[typing.Literal["query"], typing.Literal["mutation"]],
         name: str = "",
     ):
         dry_run_node = fun(PreparedArgs())
-        doc, variables, files = transport.build_gql(dry_run_node, ty, name)
+        query = (
+            {"value": dry_run_node}
+            if isinstance(dry_run_node, SelectNode)
+            else dry_run_node
+        )
+        doc, variables, files = transport.build_gql(query, ty, name)
+        self.single_node = isinstance(dry_run_node, SelectNode)
         self.doc = doc
         self._mapping = variables
         self.transport = transport
@@ -803,11 +853,14 @@ def resolve_vars(
         return resolved
 
 
-class PreparedRequestUrlib(PreparedRequestBase[Out]):
+class PreparedRequestUrlib(PreparedRequestBase[Out, PreparedOut]):
     def __init__(
         self,
         transport: GraphQLTransportUrlib,
-        fun: typing.Callable[[PreparedArgs], typing.Mapping[str, SelectNode[Out]]],
+        fun: typing.Callable[
+            [PreparedArgs],
+            typing.Union[SelectNode[Out], typing.Mapping[str, SelectNode[Out]]],
+        ],
         ty: typing.Union[typing.Literal["query"], typing.Literal["mutation"]],
         name: str = "",
     ):
@@ -818,9 +871,13 @@ def perform(
         self,
         args: typing.Mapping[str, typing.Any],
         opts: typing.Optional[GraphQLTransportOptions] = None,
-    ) -> typing.Dict[str, Out]:
+    ) -> PreparedOut:
         resolved_vars = self.resolve_vars(args, self._mapping)
-        return self.transport.fetch(self.doc, resolved_vars, opts)
+        result = self.transport.fetch(self.doc, resolved_vars, opts)
+        if self.single_node:
+            return result["value"]
+        else:
+            return result
 
 
 #

diff --git a/src/metagen/src/client_ts/static/mod.ts b/src/metagen/src/client_ts/static/mod.ts
@@ -738,7 +738,7 @@ export class GraphQLTransport {
    */
   prepareQuery<
     T extends JsonObject,
-    Q extends Record<string, QueryNode<unknown>>,
+    Q extends QueryNode<unknown> | Record<string, QueryNode<unknown>>,
   >(
     fun: (args: PreparedArgs<T>) => Q,
     { name = "" }: { name?: string } = {},
@@ -758,7 +758,7 @@ export class GraphQLTransport {
    */
   prepareMutation<
     T extends JsonObject,
-    Q extends Record<string, MutationNode<unknown>>,
+    Q extends MutationNode<unknown> | Record<string, MutationNode<unknown>>,
   >(
     fun: (args: PreparedArgs<T>) => Q,
     { name = "" }: { name?: string } = {},

diff --git a/tests/metagen/metagen_test.ts b/tests/metagen/metagen_test.ts
@@ -613,6 +613,8 @@ Meta.test(
     const expectedSchemaS = zod.object({
       singleQuery: postSchema,
       singleMutation: zod.string(),
+      singlePreparedQuery: zod.object({ input: zod.number() }),
+      singlePreparedMutation: zod.object({ input: zod.number() }),
     });
     const expectedSchema = zod.tuple([
       expectedSchemaQ,

diff --git a/tests/metagen/typegraphs/sample.ts b/tests/metagen/typegraphs/sample.ts
@@ -93,6 +93,15 @@ export const tg = await typegraph({
             list: [{ value: 3 }],
           }),
         }),
+        identity: deno.identity(t.struct({ input: t.integer() })),
+        identityUpdate: deno.func(
+          t.struct({ input: t.integer() }),
+          t.struct({ input: t.integer() }),
+          {
+            code: (input) => input,
+            effect: fx.update(),
+          },
+        ),
       },
       Policy.public(),
     );