Improve type generation accuracy

src/TypeScriptGenerator.ts

@@ -9,6 +9,7 @@ import {
   TypeGenerator,
   TypeID
 } from "relay-compiler";
+import { UnionTypeID } from "relay-compiler/lib/core/Schema";
 import { TypeGeneratorOptions } from "relay-compiler/lib/language/RelayLanguagePluginInterface";
 import * as FlattenTransform from "relay-compiler/lib/transforms/FlattenTransform";
 import * as MaskTransform from "relay-compiler/lib/transforms/MaskTransform";
@@ -127,21 +128,29 @@ function makeProp(
   selection: Selection,
   state: State,
   unmasked: boolean,
-  concreteType?: string
+  concreteType?: string,
+  unionType?: UnionTypeID
 ): ts.PropertySignature {
   let { value } = selection;
 
   const { key, schemaName, conditional, nodeType, nodeSelections } = selection;
 
   if (schemaName === "__typename" && concreteType) {
     value = ts.createLiteralTypeNode(ts.createLiteral(concreteType));
+  } else if (schemaName === "__typename" && unionType) {
+    value = ts.createUnionTypeNode(
+      schema
+        .getUnionTypes(unionType)
+        .map(type => ts.createLiteralTypeNode(ts.createLiteral(type.name)))
+    );
   } else if (nodeType) {
     value = transformScalarType(
       schema,
       nodeType,
       state,
       selectionsToAST(
         schema,
+        nodeType,
         [Array.from(nullthrows(nodeSelections).values())],
         state,
         unmasked
@@ -167,12 +176,14 @@ const onlySelectsTypename = (selections: Selection[]) =>
 
 function selectionsToAST(
   schema: Schema,
+  nodeType: TypeID | null,
   selections: ReadonlyArray<ReadonlyArray<Selection>>,
   state: State,
   unmasked: boolean,
   fragmentTypeName?: string
 ) {
   const baseFields = new Map<string, Selection>();
+  const baseFragments = new Map<string, Selection>();
 
   const byConcreteType: { [type: string]: Selection[] } = {};
 
@@ -182,6 +193,8 @@ function selectionsToAST(
     if (concreteType) {
       byConcreteType[concreteType] = byConcreteType[concreteType] || [];
       byConcreteType[concreteType].push(selection);
+    } else if (selection.ref) {
+      baseFragments.set(selection.ref, selection);
     } else {
       const previousSel = baseFields.get(selection.key);
 
@@ -205,10 +218,25 @@ function selectionsToAST(
     const typenameAliases = new Set<string>();
 
     for (const concreteType in byConcreteType) {
+      const concreteTypeSelections = byConcreteType[concreteType];
+      const concreteTypeSelectionsNames = concreteTypeSelections.map(
+        selection => selection.schemaName
+      );
+      const concreteFragmentsNames = concreteTypeSelections.map(
+        selection => selection.ref
+      );
+
       types.push(
         groupRefs([
-          ...Array.from(baseFields.values()),
-          ...byConcreteType[concreteType]
+          // Deduplicate any fields also selected on the concrete type.
+          ...Array.from(baseFields.values()).filter(
+            selection =>
+              !concreteTypeSelectionsNames.includes(selection.schemaName)
+          ),
+          ...Array.from(baseFragments.values()).filter(
+            selection => !concreteFragmentsNames.includes(selection.ref)
+          ),
+          ...concreteTypeSelections
         ]).map(selection => {
           if (selection.schemaName === "__typename") {
             typenameAliases.add(selection.key);
@@ -218,27 +246,63 @@ function selectionsToAST(
       );
     }
 
-    // It might be some other type then the listed concrete types. Ideally, we
-    // would set the type to diff(string, set of listed concrete types), but
-    // this doesn't exist in Flow at the time.
-    types.push(
-      Array.from(typenameAliases).map(typenameAlias => {
-        const otherProp = objectTypeProperty(
-          typenameAlias,
-          ts.createLiteralTypeNode(ts.createLiteral("%other"))
-        );
+    // It might be some other type then the listed concrete types. We try to
+    // figure out which types remain here.
+    let possibleTypesLeft: TypeID[] | null = null;
+    const innerType = nodeType !== null ? schema.getRawType(nodeType) : null;
+    if (innerType !== null && schema.isUnion(innerType)) {
+      const typesSeen = Object.keys(byConcreteType);
+      possibleTypesLeft = schema
+        .getUnionTypes(innerType)
+        .filter(type => !typesSeen.includes(type.name));
+    }
 
-        const otherPropWithComment = ts.addSyntheticLeadingComment(
-          otherProp,
-          ts.SyntaxKind.MultiLineCommentTrivia,
-          "This will never be '%other', but we need some\n" +
-            "value in case none of the concrete values match.",
-          true
-        );
+    // If we don't know which types are left we set the value to "%other",
+    // otherwise return a union of type names.
+    if (!possibleTypesLeft || possibleTypesLeft.length > 0) {
+      types.push([
+        ...Array.from(typenameAliases).map(typenameAlias => {
+          const otherProp = objectTypeProperty(
+            typenameAlias,
+            possibleTypesLeft
+              ? ts.createUnionTypeNode(
+                  possibleTypesLeft.map(type =>
+                    ts.createLiteralTypeNode(ts.createLiteral(type.name))
+                  )
+                )
+              : ts.createLiteralTypeNode(ts.createLiteral("%other"))
+          );
 
-        return otherPropWithComment;
-      })
-    );
+          if (possibleTypesLeft) {
+            return otherProp;
+          }
+
+          const otherPropWithComment = ts.addSyntheticLeadingComment(
+            otherProp,
+            ts.SyntaxKind.MultiLineCommentTrivia,
+            "This will never be '%other', but we need some\n" +
+              "value in case none of the concrete values match.",
+            true
+          );
+
+          return otherPropWithComment;
+        }),
+        ...(baseFragments.size > 0
+          ? objectTypeProperty(
+              FRAGMENT_REFS,
+              ts.createTypeReferenceNode(FRAGMENT_REFS_TYPE_NAME, [
+                ts.createUnionTypeNode(
+                  Array.from(baseFragments.values()).map(selection =>
+                    ts.createLiteralTypeNode(
+                      ts.createStringLiteral(selection.ref!)
+                    )
+                  )
+                )
+              ])
+            )
+          : [])
+      ]);
+    }
   } else {
     let selectionMap = selectionsToMap(Array.from(baseFields.values()));
 
@@ -254,20 +318,27 @@ function selectionsToAST(
       );
     }
 
-    const selectionMapValues = groupRefs(Array.from(selectionMap.values())).map(
-      sel =>
-        isTypenameSelection(sel) && sel.concreteType
-          ? makeProp(
-              schema,
-              {
-                ...sel,
-                conditional: false
-              },
-              state,
-              unmasked,
-              sel.concreteType
-            )
-          : makeProp(schema, sel, state, unmasked)
+    const selectionMapValues = groupRefs([
+      ...Array.from(baseFragments.values()),
+      ...Array.from(selectionMap.values())
+    ]).map(sel =>
+      isTypenameSelection(sel) &&
+      (sel.concreteType ||
+        (nodeType && schema.isUnion(schema.getNullableType(nodeType))))
+        ? makeProp(
+            schema,
+            {
+              ...sel,
+              conditional: false
+            },
+            state,
+            unmasked,
+            sel.concreteType,
+            nodeType && schema.isUnion(schema.getNullableType(nodeType))
+              ? schema.getNullableType(nodeType)
+              : undefined
+          )
+        : makeProp(schema, sel, state, unmasked)
     );
 
     types.push(selectionMapValues);
@@ -438,6 +509,7 @@ function createVisitor(
           `${node.name}Response`,
           selectionsToAST(
             schema,
+            null,
             /* $FlowFixMe: selections have already been transformed */
             (node.selections as any) as ReadonlyArray<ReadonlyArray<Selection>>,
             state,
@@ -556,6 +628,7 @@ function createVisitor(
         const unmasked = node.metadata != null && node.metadata.mask === false;
         const baseType = selectionsToAST(
           schema,
+          node.type,
           selections,
           state,
           unmasked,

test/__snapshots__/TypeScriptGenerator-test.ts.snap

@@ -446,9 +446,7 @@ export type UnionTypeTestResponse = {
         readonly __typename: "FakeNode";
         readonly id: string;
     } | {
-        /*This will never be '%other', but we need some
-        value in case none of the concrete values match.*/
-        readonly __typename: "%other";
+        readonly __typename: "NonNode";
     }) | null;
 };
 export type UnionTypeTest = {
@@ -1684,7 +1682,7 @@ export type RelayClientIDFieldQueryResponse = {
         readonly id: string;
         readonly commentBody?: {
             readonly __id: string;
-            readonly __typename: string;
+            readonly __typename: "PlainCommentBody" | "MarkdownCommentBody";
             readonly text?: {
                 readonly __id: string;
                 readonly __typename: string;
@@ -2680,9 +2678,7 @@ export type UnionTypeTestResponse = {
         readonly __typename: "FakeNode";
         readonly id: string;
     } | {
-        /*This will never be '%other', but we need some
-        value in case none of the concrete values match.*/
-        readonly __typename: "%other";
+        readonly __typename: "NonNode";
     }) | null;
 };
 export type UnionTypeTest = {
@@ -3918,7 +3914,7 @@ export type RelayClientIDFieldQueryResponse = {
         readonly id: string;
         readonly commentBody?: {
             readonly __id: string;
-            readonly __typename: string;
+            readonly __typename: "PlainCommentBody" | "MarkdownCommentBody";
             readonly text?: {
                 readonly __id: string;
                 readonly __typename: string;