Add @__RelayTypegenTypeOverride to use in CustomTransform to replace TS/JS type

compiler/crates/relay-transforms/src/lib.rs

@@ -64,6 +64,7 @@ mod skip_updatable_queries;
 mod sort_selections;
 mod test_operation_metadata;
 mod transform_connections;
+mod typegen_type_override_directive;
 mod unwrap_custom_directive_selection;
 mod util;
 mod validate_operation_variables;
@@ -203,6 +204,7 @@ pub use skip_unreachable_node::skip_unreachable_node_strict;
 pub use sort_selections::sort_selections;
 pub use test_operation_metadata::generate_test_operation_metadata;
 pub use transform_connections::transform_connections;
+pub use typegen_type_override_directive::TypgenTypeOverride;
 pub use unwrap_custom_directive_selection::unwrap_custom_directive_selection;
 pub use util::extract_variable_name;
 pub use util::generate_abstract_type_refinement_key;

compiler/crates/relay-transforms/src/typegen_type_override_directive.rs

@@ -0,0 +1,201 @@
+use std::path::PathBuf;
+
+use ::intern::string_key::Intern;
+use ::intern::string_key::StringKey;
+use common::ArgumentName;
+use common::DirectiveName;
+use common::NamedItem;
+use common::WithLocation;
+use graphql_ir::Argument;
+use graphql_ir::ConstantValue;
+use graphql_ir::Directive;
+use graphql_ir::ScalarField;
+use graphql_ir::Value;
+use intern::intern;
+use itertools::Itertools;
+use lazy_static::lazy_static;
+
+lazy_static! {
+    static ref INTERNAL_TYPE_OVERRIDE_GENERIC_ARGUMENTS: StringKey = intern!("generic_arguments");
+    static ref INTERNAL_TYPE_OVERRIDE_NAME: StringKey = intern!("name");
+    static ref INTERNAL_TYPE_OVERRIDE_PATH: StringKey = intern!("path");
+    static ref INTERNAL_TYPE_OVERRIDE_IS_UNION: StringKey = intern!("union");
+}
+
+/// The `@__RelayTypegenTypeOverride` directive allows you to specifiy a custom type
+/// that should be used in the generated code instead of the default type.
+///
+/// The directive is to be used with the `custom_transforms_config`.
+/// Specifically the `custom_transforms.typegen_transforms`
+///
+/// The directive has the following arguments:
+/// - `name`: The name of the type as it should be imported in the generated code.
+/// - `path`: The path to the file that exports the type.
+/// - `generic_arguments`: The generic arguments of the type.
+///
+/// Inside the a transform you can add this directive to a scalar using `transform_scalar_field`.
+///
+/// Example:
+/// ```
+/// use std::path::PathBuf;
+/// use std::sync::Arc;
+///
+/// use graphql_ir::ScalarField;
+/// use graphql_ir::Selection;
+/// use graphql_ir::Transformed;
+/// use graphql_ir::Transformer;
+/// use intern::string_key::Intern;
+/// use relay_transforms::TypgenTypeOverride;
+///
+/// struct MyCustomTransformer;
+/// impl Transformer for MyCustomTransformer {
+///     const NAME: &'static str = "MyCustomTransformer";
+///     const VISIT_ARGUMENTS: bool = false;
+///     const VISIT_DIRECTIVES: bool = false;
+///
+///     fn transform_scalar_field(&mut self, field: &ScalarField) -> Transformed<Selection> {
+///         let arguments = self.transform_arguments(&field.arguments);
+///         let mut directives = self
+///             .transform_directives(&field.directives)
+///             .replace_or_else(|| field.directives.clone());
+///
+///         // Check if the scalar is one that needs your custom type
+///         // By checking it's current directives or some other way
+///         let typegen_type_override = TypgenTypeOverride {
+///             name: "MyCustomScalar".intern(),
+///             path: PathBuf::from("path/to/file"),
+///             generic_arguments: vec![],
+///             is_union: false,
+///         };
+///
+///         directives.push(typegen_type_override.to_directive());
+///
+///         Transformed::Replace(Selection::ScalarField(Arc::new(ScalarField {
+///             arguments: arguments.replace_or_else(|| field.arguments.clone()),
+///             directives,
+///             ..field.clone()
+///         })))
+///     }
+/// }
+/// ```
+pub struct TypgenTypeOverride {
+    pub path: PathBuf,
+    pub name: StringKey,
+    pub generic_arguments: Vec<ConstantValue>,
+    pub is_union: bool,
+}
+
+impl TypgenTypeOverride {
+    /// Get the `@__RelayTypegenTypeOverride` directive from the scalar field
+    /// and creates a `TypgenTypeOverride` from it.
+    /// Returns None if the directive is not present
+    pub fn get_override_from_directive(scalar_field: &ScalarField) -> Option<TypgenTypeOverride> {
+        let override_typegen_type_directive_on_field =
+            TypgenTypeOverride::get_directive(scalar_field)?;
+
+        Some(TypgenTypeOverride::from(
+            override_typegen_type_directive_on_field,
+        ))
+    }
+
+    fn get_directive(scalar_field: &ScalarField) -> Option<&Directive> {
+        scalar_field.directives.named(Self::directive_name())
+    }
+
+    fn directive_name() -> DirectiveName {
+        DirectiveName("__RelayTypegenTypeOverride".intern())
+    }
+
+    pub fn to_directive(&self) -> Directive {
+        let mut arguments: Vec<Argument> = vec![
+            Argument {
+                name: WithLocation::generated(ArgumentName(*INTERNAL_TYPE_OVERRIDE_NAME)),
+                value: WithLocation::generated(Value::Constant(ConstantValue::String(self.name))),
+            },
+            Argument {
+                name: WithLocation::generated(ArgumentName(*INTERNAL_TYPE_OVERRIDE_IS_UNION)),
+                value: WithLocation::generated(Value::Constant(ConstantValue::Boolean(
+                    self.is_union,
+                ))),
+            },
+            Argument {
+                name: WithLocation::generated(ArgumentName(*INTERNAL_TYPE_OVERRIDE_PATH)),
+                value: WithLocation::generated(Value::Constant(ConstantValue::String(
+                    self.path.to_string_lossy().intern(),
+                ))),
+            },
+        ];
+
+        if !self.generic_arguments.is_empty() {
+            arguments.push(Argument {
+                name: WithLocation::generated(ArgumentName(
+                    *INTERNAL_TYPE_OVERRIDE_GENERIC_ARGUMENTS,
+                )),
+                value: WithLocation::generated(Value::Constant(ConstantValue::List(
+                    self.generic_arguments.clone(),
+                ))),
+            });
+        }
+
+        graphql_ir::Directive {
+            name: WithLocation::generated(Self::directive_name()),
+            arguments,
+            data: None,
+        }
+    }
+}
+
+impl From<&Directive> for TypgenTypeOverride {
+    fn from(directive: &Directive) -> Self {
+        let directive = directive;
+
+        let name = directive
+            .arguments
+            .named(ArgumentName(*INTERNAL_TYPE_OVERRIDE_NAME))
+            .and_then(|a| a.value.item.get_string_literal())
+            .expect("Expected a name argument for the @__RelayTypegenTypeOverride directive");
+
+        let path = directive
+            .arguments
+            .named(ArgumentName(*INTERNAL_TYPE_OVERRIDE_PATH))
+            .and_then(|a| a.value.item.get_string_literal())
+            .map(|path| PathBuf::from(path.to_string()))
+            .expect("Expected a path argument for the @__RelayTypegenTypeOverride directive");
+
+        let is_union = directive
+            .arguments
+            .named(ArgumentName(*INTERNAL_TYPE_OVERRIDE_IS_UNION))
+            .and_then(|a| match a.value.item.get_constant() {
+                Some(ConstantValue::Boolean(b)) => Some(*b),
+                _ => None,
+            })
+            .unwrap_or(false);
+
+        let generic_arguments = directive
+            .arguments
+            .named(ArgumentName(*INTERNAL_TYPE_OVERRIDE_GENERIC_ARGUMENTS))
+            .and_then(|a| match &a.value.item {
+                Value::Constant(ConstantValue::List(values)) => Some(
+                    values
+                        .iter()
+                        .filter_map(|v| match v {
+                            ConstantValue::String(s) => Some(ConstantValue::String(*s)),
+                            ConstantValue::Int(i) => Some(ConstantValue::Int(*i)),
+                            ConstantValue::Float(f) => Some(ConstantValue::Float(*f)),
+                            ConstantValue::Boolean(b) => Some(ConstantValue::Boolean(*b)),
+                            _ => None,
+                        })
+                        .collect_vec(),
+                ),
+                _ => None,
+            })
+            .unwrap_or(vec![]);
+
+        TypgenTypeOverride {
+            name,
+            path,
+            generic_arguments,
+            is_union,
+        }
+    }
+}

compiler/crates/relay-typegen/src/visit.rs

@@ -9,6 +9,7 @@ use std::collections::HashSet;
 use std::hash::Hash;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::vec;
 
 use ::intern::intern;
 use ::intern::string_key::Intern;
@@ -24,6 +25,7 @@ use docblock_shared::RELAY_RESOLVER_DIRECTIVE_NAME;
 use docblock_shared::RELAY_RESOLVER_MODEL_INSTANCE_FIELD;
 use docblock_shared::RESOLVER_VALUE_SCALAR_NAME;
 use graphql_ir::Condition;
+use graphql_ir::ConstantValue;
 use graphql_ir::Directive;
 use graphql_ir::FragmentDefinitionName;
 use graphql_ir::FragmentSpread;
@@ -55,6 +57,7 @@ use relay_transforms::RelayResolverMetadata;
 use relay_transforms::RequiredMetadataDirective;
 use relay_transforms::ResolverOutputTypeInfo;
 use relay_transforms::TypeConditionInfo;
+use relay_transforms::TypgenTypeOverride;
 use relay_transforms::ASSIGNABLE_DIRECTIVE_FOR_TYPEGEN;
 use relay_transforms::CHILDREN_CAN_BUBBLE_METADATA_KEY;
 use relay_transforms::CLIENT_EXTENSION_DIRECTIVE_NAME;
@@ -1202,7 +1205,9 @@ fn visit_scalar_field(
     }
 
     let ast = transform_type_reference_into_ast(&field_type, |type_| {
-        expect_scalar_type(typegen_context, encountered_enums, custom_scalars, type_)
+        override_typegen_type_directive(scalar_field, custom_scalars).unwrap_or_else(|| {
+            expect_scalar_type(typegen_context, encountered_enums, custom_scalars, type_)
+        })
     });
 
     type_selections.push(TypeSelection::ScalarField(TypeSelectionScalarField {
@@ -1215,6 +1220,44 @@ fn visit_scalar_field(
     }));
 }
 
+fn override_typegen_type_directive(
+    scalar_field: &ScalarField,
+    custom_scalars: &mut CustomScalarsImports,
+) -> Option<AST> {
+    let TypgenTypeOverride {
+        name,
+        path,
+        generic_arguments,
+        is_union,
+    } = TypgenTypeOverride::get_override_from_directive(scalar_field)?;
+
+    custom_scalars.insert((name, path.clone()));
+
+    // Override without arguments
+    if generic_arguments.is_empty() {
+        return Some(AST::RawType(name));
+    }
+
+    let arguments = generic_arguments
+        .iter()
+        .filter_map(|v| match v {
+            ConstantValue::String(s) => Some(AST::StringLiteral(StringLiteral(*s))),
+            ConstantValue::Int(i) => Some(AST::RawType(format!("{}", i).intern())),
+            ConstantValue::Float(f) => Some(AST::RawType(format!("{}", f).intern())),
+            ConstantValue::Boolean(b) => Some(AST::RawType(format!("{}", b).intern())),
+            _ => panic!("Unsupported type for typegen override argument. Only String, Int, Float and Boolean are supported."),
+        })
+        .collect();
+
+    Some(AST::GenericType {
+        outer: name,
+        inner: match is_union {
+            true => vec![AST::Union(SortedASTList::new(arguments))],
+            false => arguments,
+        },
+    })
+}
+
 #[allow(clippy::too_many_arguments)]
 fn visit_condition(
     typegen_context: &'_ TypegenContext<'_>,

compiler/crates/relay-typegen/tests/generate_flow/fixtures/relay_custom_scalar_type_boolean_argument.expected

@@ -0,0 +1,39 @@
+==================================== INPUT ====================================
+fragment foo on ViewData {
+    number @__RelayTypegenTypeOverride(
+      path: "isEven",
+      name: "IsEven",
+      generic_arguments: [
+        true
+      ]
+    )
+}
+
+# %extensions%
+
+extend type Query {
+  viewData: ViewData
+}
+
+type ViewData {
+  number: Float
+}
+
+directive @__RelayTypegenTypeOverride(
+  path: String!, 
+  name: String!,
+  generic_arguments: [Boolean!]
+) on FIELD
+==================================== OUTPUT ===================================
+import type { IsEven } from "isEven";
+import type { FragmentType } from "relay-runtime";
+declare export opaque type foo$fragmentType: FragmentType;
+export type foo$data = {|
+  +number: ?IsEven<true>,
+  +$fragmentType: foo$fragmentType,
+|};
+export type foo$key = {
+  +$data?: foo$data,
+  +$fragmentSpreads: foo$fragmentType,
+  ...
+};

compiler/crates/relay-typegen/tests/generate_flow/fixtures/relay_custom_scalar_type_boolean_argument.graphql

@@ -0,0 +1,25 @@
+fragment foo on ViewData {
+    number @__RelayTypegenTypeOverride(
+      path: "isEven",
+      name: "IsEven",
+      generic_arguments: [
+        true
+      ]
+    )
+}
+
+# %extensions%
+
+extend type Query {
+  viewData: ViewData
+}
+
+type ViewData {
+  number: Float
+}
+
+directive @__RelayTypegenTypeOverride(
+  path: String!, 
+  name: String!,
+  generic_arguments: [Boolean!]
+) on FIELD