Add CLI flag to override typeshed

pyrefly/lib/commands/check.rs

@@ -174,6 +174,9 @@ struct ConfigOverrideArgs {
     /// after first checking `search_path` and `typeshed`.
     #[arg(long, env = clap_env("SITE_PACKAGE_PATH"))]
     site_package_path: Option<Vec<PathBuf>>,
+    /// Override the bundled typeshed with a custom path.
+    #[arg(long, env = clap_env("TYPESHED_PATH"))]
+    typeshed_path: Option<PathBuf>,
     /// The Python executable that will be queried for `python_version`
     /// `python_platform`, or `site_package_path` if any of the values are missing.
     #[arg(long, env = clap_env("PYTHON_INTERPRETER"), value_name = "EXE_PATH")]
@@ -555,6 +558,9 @@ impl Args {
             config.python_environment.site_package_path = Some(x.clone());
             config.python_environment.site_package_path_source = SitePackagePathSource::CommandLine;
         }
+        if let Some(x) = &self.config_override.typeshed_path {
+            config.typeshed_path = Some(x.clone());
+        }
         if let Some(x) = &self.config_override.python_interpreter {
             config.python_interpreter = Some(x.clone());
         }

pyrefly/lib/config/config.rs

@@ -223,6 +223,13 @@ pub struct ConfigFile {
     )]
     pub fallback_search_path: Vec<PathBuf>,
 
+    /// Override the bundled typeshed with a custom path.
+    #[serde(
+        default,
+        skip_serializing_if = "Option::is_none",
+    )]
+    pub typeshed_path: Option<PathBuf>,
+
     // TODO(connernilsen): make this mutually exclusive with venv/conda env
     /// The python executable that will be queried for `python_version`,
     /// `python_platform`, or `site_package_path` if any of the values are missing.
@@ -302,6 +309,7 @@ impl Default for ConfigFile {
             custom_module_paths: Default::default(),
             use_untyped_imports: true,
             ignore_missing_source: true,
+            typeshed_path: None,
         }
     }
 }
@@ -354,14 +362,28 @@ impl ConfigFile {
             Err(FindError::Ignored)
         } else if let Some(path) = find_module_in_search_path(module, self.search_path()) {
             Ok(path)
-        } else if let Some(path) = typeshed()
+        } else if let Some(custom_typeshed_path) = &self.typeshed_path {
+            // Check custom typeshed stdlib path before bundled typeshed
+            let stdlib_path = custom_typeshed_path.join("stdlib");
+            if let Some(path) = find_module_in_search_path(module, std::iter::once(&stdlib_path)) {
+                Ok(path)
+            } else {
+                // If not found in custom typeshed, continue with normal lookup
+                self.continue_import_lookup(module)
+            }
+        } else if self.typeshed_path.is_none() && let Some(path) = typeshed()
             .map_err(|err| FindError::not_found(err, module))?
             .find(module)
         {
             Ok(path)
-        } else if let Some(path) =
-            find_module_in_search_path(module, self.fallback_search_path.iter())
-        {
+        } else {
+            // Continue with normal lookup
+            self.continue_import_lookup(module)
+        }
+    }
+
+    fn continue_import_lookup(&self, module: ModuleName) -> Result<ModulePath, FindError> {
+        if let Some(path) = find_module_in_search_path(module, self.fallback_search_path.iter()) {
             Ok(path)
         } else if let Some(path) = find_module_in_site_package_path(
             module,
@@ -823,6 +845,7 @@ mod tests {
                 }],
                 use_untyped_imports: true,
                 ignore_missing_source: true,
+                typeshed_path: None,
             }
         );
     }
@@ -1033,6 +1056,7 @@ mod tests {
             }],
             use_untyped_imports: false,
             ignore_missing_source: false,
+            typeshed_path: None,
         };
 
         let path_str = with_sep("path/to/my/config");
@@ -1069,6 +1093,7 @@ mod tests {
             }],
             use_untyped_imports: false,
             ignore_missing_source: false,
+            typeshed_path: None,
         };
         assert_eq!(config, expected_config);
     }