From 97ade9fa46f4436122e7b6a94287546a0319d62b Mon Sep 17 00:00:00 2001
From: izumin5210 <m@izum.in>
Date: Mon, 13 Sep 2021 22:50:14 +0900
Subject: [PATCH] Return null when input is null

---
 src/maskWithFragment.test.ts | 6 ++++++
 src/maskWithFragment.ts      | 8 ++++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/maskWithFragment.test.ts b/src/maskWithFragment.test.ts
index c4d40b3..448694d 100644
--- a/src/maskWithFragment.test.ts
+++ b/src/maskWithFragment.test.ts
@@ -8,6 +8,12 @@ import * as multipleFragmentsFixtures from "./__fixtures__/graphql/__generated__
 import * as rootListFixtures from "./__fixtures__/graphql/__generated__/rootList.generated";
 import * as unionFixtures from "./__fixtures__/graphql/__generated__/union.generated";
 
+it("returns null when 2nd arg is null", () => {
+  const output = maskWithFragment(simpleFixtures.UserHeaderFragmentDoc, null);
+  expectType<null>(output);
+  expect(output).toBeNull();
+});
+
 it("masks query results with simple fragment", () => {
   const input: simpleFixtures.GetUserHeaderQuery = {
     __typename: "Query",
diff --git a/src/maskWithFragment.ts b/src/maskWithFragment.ts
index 2840600..4a53e7f 100644
--- a/src/maskWithFragment.ts
+++ b/src/maskWithFragment.ts
@@ -82,13 +82,17 @@ export function maskWithFragment<
   TFilteredData extends Record<string, unknown>,
   TData extends TFilteredData = TFilteredData
 >(doc: DocumentNode | TypedDocumentNode<TFilteredData, any>, input: ReadonlyArray<TData>): ReadonlyArray<TFilteredData>;
+export function maskWithFragment(doc: DocumentNode | TypedDocumentNode<any, any>, input: null): null;
+export function maskWithFragment(doc: DocumentNode | TypedDocumentNode<any, any>, input: undefined): undefined;
 export function maskWithFragment<
   TFilteredData extends Record<string, unknown>,
   TData extends TFilteredData = TFilteredData
 >(
   doc: DocumentNode | TypedDocumentNode<TFilteredData, any>,
-  input: TData | ReadonlyArray<TData>
-): TFilteredData | ReadonlyArray<TFilteredData> {
+  input: TData | ReadonlyArray<TData> | null | undefined
+): TFilteredData | ReadonlyArray<TFilteredData> | null | undefined {
+  if (input == null) return input;
+
   if (doc.definitions[0]?.kind !== "FragmentDefinition") {
     throw new Error("input document should be fragment definition");
   }