Make typeshed_path a full replacement for bundled stubs

pyrefly/lib/lsp/non_wasm/workspace.rs

@@ -85,6 +85,11 @@ pub struct Workspace {
     pub stream_diagnostics: Option<bool>,
     pub diagnostic_mode: Option<DiagnosticMode>,
     pub workspace_config: Option<PathBuf>,
+    /// Custom typeshed override supplied over LSP. Equivalent to the
+    /// `typeshed_path` field of a `pyrefly.toml`, but settable via
+    /// `initializationOptions`/`workspace/configuration` so editor integrations
+    /// can point Pyrefly at their own typeshed without writing a config file.
+    pub typeshed_path: Option<PathBuf>,
 }
 
 impl Workspace {
@@ -172,6 +177,13 @@ impl ConfigConfigurer for WorkspaceConfigConfigurer {
                     // skip interpreter query because we already have the interpreter from the workspace
                     config.interpreters.skip_interpreter_query = true;
                 }
+                // Workspace-level `typeshed_path` from LSP overrides the bundled typeshed
+                // but defers to an explicit value already set by a config file the user owns.
+                if let Some(typeshed_path) = w.typeshed_path.clone()
+                    && config.typeshed_path.is_none()
+                {
+                    config.typeshed_path = Some(typeshed_path);
+                }
             })
         };
 
@@ -252,6 +264,10 @@ struct PyreflyClientConfig {
     disabled_language_services: Option<DisabledLanguageServices>,
     stream_diagnostics: Option<bool>,
     config_path: Option<PathBuf>,
+    /// Override the bundled typeshed with a custom path. Equivalent to the
+    /// `typeshed_path` field of a `pyrefly.toml`, but settable over the LSP
+    /// `initializationOptions`/`workspace/configuration` channel.
+    typeshed_path: Option<PathBuf>,
 }
 
 #[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq)]
@@ -544,6 +560,7 @@ impl Workspaces {
             // so a partial `did_change_configuration` touching some other
             // key won't pay for a needless recheck.
             self.update_display_type_errors(modified, scope_uri, pyrefly.display_type_errors);
+            self.update_typeshed_path(modified, scope_uri, pyrefly.typeshed_path);
             self.update_type_checking_mode(
                 modified,
                 scope_uri,
@@ -656,6 +673,37 @@ impl Workspaces {
     }
 
     /// Update displayTypeErrors setting for scope_uri, None if default workspace
+    /// Update [`Workspace::typeshed_path`] for `scope_uri`, or the default workspace
+    /// when `scope_uri` is `None`. `modified` is set only when the value actually
+    /// changes, so a partial payload re-stating the current value doesn't trigger a
+    /// spurious recheck.
+    fn update_typeshed_path(
+        &self,
+        modified: &mut bool,
+        scope_uri: &Option<Url>,
+        typeshed_path: Option<PathBuf>,
+    ) {
+        let mut workspaces = self.workspaces.write();
+        match scope_uri {
+            Some(scope_uri) => {
+                if let Ok(path) = scope_uri.to_file_path()
+                    && let Some(workspace) = workspaces.get_mut(&path)
+                    && workspace.typeshed_path != typeshed_path
+                {
+                    *modified = true;
+                    workspace.typeshed_path = typeshed_path;
+                }
+            }
+            None => {
+                let mut default = self.default.write();
+                if default.typeshed_path != typeshed_path {
+                    *modified = true;
+                    default.typeshed_path = typeshed_path;
+                }
+            }
+        }
+    }
+
     fn update_display_type_errors(
         &self,
         modified: &mut bool,
@@ -1093,6 +1141,37 @@ mod tests {
         assert!(config.pyrefly.is_some());
     }
 
+    #[test]
+    fn test_pyrefly_typeshed_path_parsed_from_lsp_config() {
+        let lsp_config: LspConfig = serde_json::from_value(json!({
+            "pyrefly": { "typeshedPath": "/path/to/custom/typeshed" }
+        }))
+        .expect("LspConfig should parse");
+        assert_eq!(
+            lsp_config.pyrefly.expect("pyrefly section").typeshed_path,
+            Some(PathBuf::from("/path/to/custom/typeshed"))
+        );
+    }
+
+    #[test]
+    fn test_pyrefly_typeshed_path_applied_to_default_workspace() {
+        let workspaces = Workspaces::new(Workspace::new(), &[]);
+        let mut modified = false;
+        workspaces.apply_client_configuration(
+            &mut modified,
+            &None,
+            json!({ "pyrefly": { "typeshedPath": "/path/to/custom/typeshed" } }),
+        );
+        assert!(
+            modified,
+            "applying typeshed_path should mark workspace modified"
+        );
+        assert_eq!(
+            workspaces.default.read().typeshed_path,
+            Some(PathBuf::from("/path/to/custom/typeshed"))
+        );
+    }
+
     /// Legacy `displayTypeErrors` maps onto the two new axes:
     /// `force-on` sets the typeCheckingMode to Default and is a no-op
     /// on the kill switch; `force-off` sets the kill switch and is a

pyrefly/lib/module/finder.rs

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::collections::HashMap;
 use std::ffi::OsString;
 use std::fmt::Debug;
 use std::io::Read;
@@ -20,6 +21,7 @@ use pyrefly_python::COMPILED_FILE_SUFFIXES;
 use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::module_path::ModulePath;
 use pyrefly_python::module_path::ModuleStyle;
+use pyrefly_util::lock::Mutex;
 use pyrefly_util::locked_map::LockedMap;
 use pyrefly_util::suggest::best_suggestion;
 use regex::Regex;
@@ -1115,14 +1117,79 @@ fn find_extra_extension_module<'a>(
     None
 }
 
+/// Per-typeshed cache of the immediate subdirectories of `<typeshed>/stubs`. Each
+/// subdirectory is a typeshed-style stub distribution (`stubs/Markdown/`,
+/// `stubs/pandas-stubs/`, ...), and we treat each of them as its own search root,
+/// mirroring typeshed's third-party layout.
+static CUSTOM_TYPESHED_STUB_PKG_DIRS: LazyLock<Mutex<HashMap<PathBuf, Arc<Vec<PathBuf>>>>> =
+    LazyLock::new(|| Mutex::new(HashMap::new()));
+
+fn custom_typeshed_stub_pkg_dirs(typeshed_root: &Path) -> Arc<Vec<PathBuf>> {
+    let stubs_root = typeshed_root.join("stubs");
+    let mut cache = CUSTOM_TYPESHED_STUB_PKG_DIRS.lock();
+    if let Some(cached) = cache.get(&stubs_root) {
+        return cached.clone();
+    }
+    let pkgs: Vec<PathBuf> = std::fs::read_dir(&stubs_root)
+        .ok()
+        .into_iter()
+        .flatten()
+        .filter_map(Result::ok)
+        .filter_map(|entry| {
+            let path = entry.path();
+            if path.is_dir() { Some(path) } else { None }
+        })
+        .collect();
+    let arc = Arc::new(pkgs);
+    cache.insert(stubs_root, arc.clone());
+    arc
+}
+
+fn find_in_custom_typeshed_stubs(
+    typeshed_root: &Path,
+    module: ModuleName,
+    style_filter: Option<ModuleStyle>,
+) -> Option<FindingOrError<ModulePath>> {
+    let pkg_dirs = custom_typeshed_stub_pkg_dirs(typeshed_root);
+    if pkg_dirs.is_empty() {
+        return None;
+    }
+    // Namespaces and phantom-paths discovered while scanning typeshed stubs are
+    // discarded: typeshed stub packages aren't expected to form namespace packages
+    // and they are not surfaced to the outer `find_import_internal` loop anyway.
+    let mut namespaces_found = Vec::new();
+    let mut phantom_paths: Option<&mut Vec<PathBuf>> = None;
+    let dir_cache = DirEntryCache::new(false);
+    find_module(
+        module,
+        pkg_dirs.iter(),
+        &mut namespaces_found,
+        style_filter,
+        None,
+        false,
+        &mut phantom_paths,
+        &dir_cache,
+        None,
+    )
+}
+
 /// This function will find either third party typeshed stubs or other third party stubs
 /// Here a decision is being made to prioritize typeshed stubs over other third party stubs that are bundled.
 /// Since we run the typeshed update script with a more regular cadence, it is more likely that
 /// these stubs will be more up to date.
 fn find_third_party_stub(
+    config: &ConfigFile,
     module: ModuleName,
     style_filter: Option<ModuleStyle>,
 ) -> Option<FindingOrError<ModulePath>> {
+    // A configured `typeshed_path` fully replaces Pyrefly's bundled stubs: third-party
+    // stubs are resolved only from `<typeshed_path>/stubs`, and a miss returns `None`
+    // so that later `find_import_internal` stages (site-packages, ...) take over. We
+    // never fall back to the bundled typeshed/Meta stubs.
+    if let Some(custom_typeshed_path) = &config.typeshed_path {
+        return find_in_custom_typeshed_stubs(custom_typeshed_path, module, style_filter);
+    }
+
     let third_party_typeshed_stub = if matches!(style_filter, Some(ModuleStyle::Interface) | None) {
         typeshed_third_party().map_or_else(
             |err| {
@@ -1187,7 +1254,7 @@ pub fn find_import_internal(
 ) -> FindingOrError<ModulePath> {
     let mut namespaces_found = vec![];
     let origin = origin.map(|p| p.as_path());
-    let typeshed_third_party_result = find_third_party_stub(module, style_filter);
+    let typeshed_third_party_result = find_third_party_stub(config, module, style_filter);
     let typeshed_third_party_stub = typeshed_third_party_result.clone();
     let from_real_config_file = config.from_real_config_file();
 
@@ -1237,7 +1304,8 @@ pub fn find_import_internal(
         )
     {
         path
-    } else if matches!(style_filter, Some(ModuleStyle::Interface) | None)
+    } else if config.typeshed_path.is_none()
+        && matches!(style_filter, Some(ModuleStyle::Interface) | None)
         && let Some(path) = typeshed().map_or_else(
             |err| {
                 Some(FindingOrError::Error(FindError::missing_import(
@@ -1247,6 +1315,9 @@ pub fn find_import_internal(
             |ts| ts.find(module).map(FindingOrError::new_finding),
         )
     {
+        // Skip the bundled typeshed stdlib entirely when the user pointed us at their
+        // own typeshed: a configured `typeshed_path` is consulted above and must never
+        // fall back to Pyrefly's bundled stubs.
         path
     } else if !config.disable_search_path_heuristics
         && let Some(path) = find_module(
@@ -4262,7 +4333,8 @@ mod tests {
     fn test_find_third_party_stub_prioritizes_typeshed_over_bundled() {
         // 'requests' exists in typeshed third party stubs, so it should
         // return BundledTypeshedThirdParty (typeshed is prioritized over other bundled stubs)
-        let result = find_third_party_stub(ModuleName::from_str("requests"), None);
+        let config = get_config(ConfigSource::Synthetic);
+        let result = find_third_party_stub(&config, ModuleName::from_str("requests"), None);
         assert!(result.is_some(), "Should find 'requests' stub");
 
         if let Some(FindingOrError::Finding(finding)) = result {
@@ -4283,7 +4355,8 @@ mod tests {
     fn test_find_third_party_stub_returns_bundled_when_not_in_typeshed() {
         // 'conans' exists only in BundledThirdParty (from third_party/stubs/conans-stubs),
         // not in typeshed, so it should return BundledThirdParty
-        let result = find_third_party_stub(ModuleName::from_str("conans"), None);
+        let config = get_config(ConfigSource::Synthetic);
+        let result = find_third_party_stub(&config, ModuleName::from_str("conans"), None);
         assert!(result.is_some(), "Should find 'conans' stub");
 
         if let Some(FindingOrError::Finding(finding)) = result {
@@ -4300,6 +4373,62 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_find_third_party_stub_skips_bundled_when_typeshed_path_set() {
+        // `conans` (Meta-bundled) and `requests` (bundled typeshed) both resolve via
+        // the bundled stubs by default. Once `typeshed_path` points at an empty custom
+        // typeshed, neither bundled source is consulted, so both must miss.
+        let mut config = get_config(ConfigSource::Synthetic);
+        assert!(
+            find_third_party_stub(&config, ModuleName::from_str("conans"), None).is_some(),
+            "Sanity check: 'conans' resolves via bundled by default"
+        );
+        assert!(
+            find_third_party_stub(&config, ModuleName::from_str("requests"), None).is_some(),
+            "Sanity check: 'requests' resolves via bundled typeshed by default"
+        );
+
+        let tmp = tempfile::tempdir().expect("create tempdir");
+        std::fs::create_dir_all(tmp.path().join("stubs")).expect("create empty stubs dir");
+        config.typeshed_path = Some(tmp.path().to_path_buf());
+        assert!(
+            find_third_party_stub(&config, ModuleName::from_str("conans"), None).is_none(),
+            "Bundled Meta stubs must be skipped when typeshed_path is set"
+        );
+        assert!(
+            find_third_party_stub(&config, ModuleName::from_str("requests"), None).is_none(),
+            "Bundled typeshed third-party stubs must also be skipped when typeshed_path is set"
+        );
+    }
+
+    #[test]
+    fn test_find_third_party_stub_prefers_custom_typeshed() {
+        // Build a typeshed-style directory on disk with one stub package and verify
+        // that `config.typeshed_path` makes Pyrefly resolve modules out of it instead
+        // of falling through to the bundled typeshed/third-party stubs.
+        let tmp = tempfile::tempdir().expect("create tempdir");
+        let stubs_root = tmp.path().join("stubs");
+        let pkg = stubs_root.join("custom_pkg");
+        std::fs::create_dir_all(pkg.join("custom_pkg")).expect("create pkg dirs");
+        std::fs::write(pkg.join("custom_pkg/__init__.pyi"), "X: int = 0\n").expect("write stub");
+
+        let mut config = get_config(ConfigSource::Synthetic);
+        config.typeshed_path = Some(tmp.path().to_path_buf());
+
+        let result = find_third_party_stub(&config, ModuleName::from_str("custom_pkg"), None)
+            .expect("custom typeshed lookup hits");
+        match result {
+            FindingOrError::Finding(finding) => {
+                assert!(
+                    matches!(finding.finding.details(), ModulePathDetails::FileSystem(_)),
+                    "Expected FileSystem path for custom typeshed, got: {:?}",
+                    finding.finding.details()
+                );
+            }
+            other => panic!("Expected Finding from custom typeshed, got: {:?}", other),
+        }
+    }
+
     #[test]
     fn test_suggest_stdlib_import() {
         // Test that we suggest 'math' for 'mathh' (one character typo)

pyrefly/lib/module/typeshed.rs

@@ -78,21 +78,49 @@ impl BundledStub for BundledTypeshedStdlib {
                 Some(path) => vec![path],
                 None => Vec::new(),
             };
-            let error_overrides = HashMap::from([
-                // The stdlib is full of deliberately incorrect overrides, so ignore them
-                (ErrorKind::BadOverride, Severity::Ignore),
-                (ErrorKind::BadOverrideParamName, Severity::Ignore),
-                // The stdlib has variance violations in typing.pyi, so ignore them
-                (ErrorKind::InvalidVariance, Severity::Ignore),
-            ]);
-            let config_file =
-                create_bundled_stub_config(Some(search_paths), Some(error_overrides), Some(true));
+            let config_file = create_bundled_stub_config(
+                Some(search_paths),
+                Some(stdlib_error_overrides()),
+                Some(true),
+            );
             ArcId::new(config_file)
         });
         CONFIG.dupe()
     }
 }
 
+/// The stdlib is full of deliberately incorrect overrides and variance violations
+/// (e.g. in `typing.pyi`), so silence those error kinds when type-checking it.
+fn stdlib_error_overrides() -> HashMap<ErrorKind, Severity> {
+    HashMap::from([
+        (ErrorKind::BadOverride, Severity::Ignore),
+        (ErrorKind::BadOverrideParamName, Severity::Ignore),
+        (ErrorKind::InvalidVariance, Severity::Ignore),
+    ])
+}
+
+/// Loader config used to resolve and type-check the stdlib.
+///
+/// With no `typeshed_path`, the bundled stdlib snapshot is used. When a
+/// `typeshed_path` is configured, the returned config carries it so that the
+/// normal finder branch resolves stdlib modules out of `<typeshed_path>/stdlib`
+/// and never falls back to the bundled snapshot — making a configured typeshed a
+/// complete replacement for Pyrefly's bundled stdlib, just like third-party stubs.
+pub fn stdlib_config(typeshed_path: Option<&Path>) -> ArcId<ConfigFile> {
+    match typeshed_path {
+        None => BundledTypeshedStdlib::config(),
+        Some(typeshed_path) => {
+            // `typeshed_path` is already absolute (read from an
+            // already-configured `ConfigFile`), so it needs no further
+            // absolutization after `create_bundled_stub_config` configures.
+            let mut config_file =
+                create_bundled_stub_config(None, Some(stdlib_error_overrides()), Some(true));
+            config_file.typeshed_path = Some(typeshed_path.to_path_buf());
+            ArcId::new(config_file)
+        }
+    }
+}
+
 static BUNDLED_TYPESHED: LazyLock<anyhow::Result<BundledTypeshedStdlib>> =
     LazyLock::new(BundledTypeshedStdlib::new);