ty_module_resolver/
module.rs

1use std::borrow::Cow;
2use std::fmt::Formatter;
3use std::str::FromStr;
4
5use ruff_db::files::{File, directory_listing, system_path_to_file, vendored_path_to_file};
6use ruff_db::system::SystemPath;
7use ruff_db::vendored::VendoredPath;
8use ruff_python_ast::PythonVersion;
9use salsa::Database;
10use salsa::plumbing::AsId;
11
12use crate::module_name::ModuleName;
13use crate::path::{SearchPath, SystemOrVendoredPathRef};
14use crate::{Db, ResolverEnvironment};
15
16/// Representation of a Python module.
17#[derive(Clone, Copy, Eq, Hash, PartialEq, salsa::Supertype, salsa::SalsaValue)]
18pub enum Module<'db> {
19    File(FileModule<'db>),
20    Namespace(NamespacePackage<'db>),
21}
22
23// The Salsa heap is tracked separately.
24impl get_size2::GetSize for Module<'_> {}
25
26#[salsa::tracked]
27impl<'db> Module<'db> {
28    pub(crate) fn file_module(
29        db: &'db dyn Db,
30        file: File,
31        resolver_environment: ResolverEnvironment<'db>,
32        name: Cow<'_, ModuleName>,
33        kind: ModuleKind,
34        search_path: SearchPath,
35    ) -> Self {
36        let known = KnownModule::try_from_search_path_and_name(&search_path, &name);
37
38        Self::File(FileModule::new(
39            db,
40            name,
41            kind,
42            search_path,
43            file,
44            resolver_environment,
45            known,
46        ))
47    }
48
49    pub(crate) fn namespace_package(
50        db: &'db dyn Db,
51        resolver_environment: ResolverEnvironment<'db>,
52        name: Cow<'_, ModuleName>,
53    ) -> Self {
54        Self::Namespace(NamespacePackage::new(db, resolver_environment, name))
55    }
56
57    /// The resolver environment used to resolve this module.
58    pub fn resolver_environment(self, db: &'db dyn Database) -> ResolverEnvironment<'db> {
59        match self {
60            Module::File(module) => module.resolver_environment(db),
61            Module::Namespace(module) => module.resolver_environment(db),
62        }
63    }
64
65    /// The absolute name of the module (e.g. `foo.bar`)
66    pub fn name(self, db: &'db dyn Database) -> &'db ModuleName {
67        match self {
68            Module::File(module) => module.name(db),
69            Module::Namespace(ref package) => package.name(db),
70        }
71    }
72
73    /// The file to the source code that defines this module
74    ///
75    /// This is `None` for namespace packages.
76    pub fn file(self, db: &'db dyn Database) -> Option<File> {
77        match self {
78            Module::File(module) => Some(module.file(db)),
79            Module::Namespace(_) => None,
80        }
81    }
82
83    /// The Python version used to resolve this module.
84    pub fn python_version(self, db: &'db dyn Database) -> PythonVersion {
85        self.resolver_environment(db).python_version(db)
86    }
87
88    /// Is this a module that we special-case somehow? If so, which one?
89    pub fn known(self, db: &'db dyn Database) -> Option<KnownModule> {
90        match self {
91            Module::File(module) => module.known(db),
92            Module::Namespace(_) => None,
93        }
94    }
95
96    /// Does this module represent the given known module?
97    pub fn is_known(self, db: &'db dyn Database, known_module: KnownModule) -> bool {
98        self.known(db) == Some(known_module)
99    }
100
101    /// The search path from which the module was resolved.
102    ///
103    /// It is guaranteed that if `None` is returned, then this is a namespace
104    /// package. Otherwise, this is a regular package or file module.
105    pub fn search_path(self, db: &'db dyn Database) -> Option<&'db SearchPath> {
106        match self {
107            Module::File(module) => Some(module.search_path(db)),
108            Module::Namespace(_) => None,
109        }
110    }
111
112    /// Returns whether this module resolves to a bundled typing-only stub.
113    ///
114    /// A project or installed module with the same name may still exist on a
115    /// lower-priority search path and be available at runtime.
116    pub fn is_type_check_only(self, db: &'db dyn Database) -> bool {
117        self.search_path(db)
118            .is_some_and(SearchPath::is_standard_library)
119            && matches!(
120                self.name(db).first_component(),
121                "_typeshed" | "typing_extensions" | "ty_extensions"
122            )
123    }
124
125    /// Determine whether this module is a single-file module or a package
126    pub fn kind(self, db: &'db dyn Database) -> ModuleKind {
127        match self {
128            Module::File(module) => module.kind(db),
129            Module::Namespace(_) => ModuleKind::Package,
130        }
131    }
132
133    /// Return a list of all submodules of this module.
134    ///
135    /// Returns an empty list if the module is not a package, if it is an empty package,
136    /// or if it is a namespace package (one without an `__init__.py` or `__init__.pyi` file).
137    ///
138    /// The names returned correspond to the "base" name of the module.
139    /// That is, `{self.name}.{basename}` should give the full module name.
140    pub fn all_submodules(self, db: &'db dyn Db) -> &'db [Module<'db>] {
141        all_submodule_names_for_package(db, self).unwrap_or_default()
142    }
143}
144
145impl std::fmt::Debug for Module<'_> {
146    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
147        salsa::with_attached_database(|db| {
148            f.debug_struct("Module")
149                .field("name", &self.name(db))
150                .field("kind", &self.kind(db))
151                .field("file", &self.file(db))
152                .field("search_path", &self.search_path(db))
153                .field("known", &self.known(db))
154                .finish()
155        })
156        .unwrap_or_else(|| f.debug_tuple("Module").field(&self.as_id()).finish())
157    }
158}
159
160#[salsa::tracked(returns(as_deref), heap_size=ruff_memory_usage::heap_size)]
161fn all_submodule_names_for_package<'db>(
162    db: &'db dyn Db,
163    module: Module<'db>,
164) -> Option<Box<[Module<'db>]>> {
165    fn is_submodule(
166        is_dir: bool,
167        is_file: bool,
168        basename: Option<&str>,
169        extension: Option<&str>,
170    ) -> bool {
171        is_dir
172            || (is_file
173                && matches!(extension, Some("py" | "pyi"))
174                && !matches!(basename, Some("__init__.py" | "__init__.pyi")))
175    }
176
177    fn find_package_init_system(db: &dyn Db, dir: &SystemPath) -> Option<File> {
178        let listing = directory_listing(db, dir).ok()?;
179        if listing.entry_is_file(db, dir, "__init__.pyi") {
180            system_path_to_file(db, dir.join("__init__.pyi")).ok()
181        } else if listing.entry_is_file(db, dir, "__init__.py") {
182            system_path_to_file(db, dir.join("__init__.py")).ok()
183        } else {
184            None
185        }
186    }
187
188    fn find_package_init_vendored(db: &dyn Db, dir: &VendoredPath) -> Option<File> {
189        vendored_path_to_file(db, dir.join("__init__.pyi"))
190            .or_else(|_| vendored_path_to_file(db, dir.join("__init__.py")))
191            .ok()
192    }
193
194    // It would be complex and expensive to compute all submodules for
195    // namespace packages, since a namespace package doesn't correspond
196    // to a single file; it can span multiple directories across multiple
197    // search paths. For now, we only compute submodules for traditional
198    // packages that exist in a single directory on a single search path.
199    let Module::File(module) = module else {
200        return None;
201    };
202    if !matches!(module.kind(db), ModuleKind::Package) {
203        return None;
204    }
205
206    let path = SystemOrVendoredPathRef::try_from_file(db, module.file(db))?;
207    debug_assert!(
208        matches!(path.file_name(), Some("__init__.py" | "__init__.pyi")),
209        "expected package file `{:?}` to be `__init__.py` or `__init__.pyi`",
210        path.file_name(),
211    );
212
213    let resolver_environment = module.resolver_environment(db);
214    Some(match path.parent()? {
215        SystemOrVendoredPathRef::System(parent_directory) => {
216            directory_listing(db, parent_directory)
217                .inspect_err(|error| {
218                    tracing::debug!(
219                        "Failed to read {parent_directory:?} when looking for \
220                         its possible submodules: {error}"
221                    );
222                })
223                .ok()?
224                .iter()
225                .filter(|(name, ty)| {
226                    let path = SystemPath::new(name);
227                    is_submodule(
228                        ty.is_directory(),
229                        ty.is_file(),
230                        path.file_name(),
231                        path.extension(),
232                    )
233                })
234                .filter_map(|(entry_name, file_type)| {
235                    let relative = SystemPath::new(entry_name);
236                    let stem = relative.file_stem()?;
237                    let path = parent_directory.join(relative);
238                    let mut name = module.name(db).clone();
239                    name.extend(&ModuleName::new(stem)?);
240
241                    let (kind, file) = if file_type.is_directory() {
242                        (ModuleKind::Package, find_package_init_system(db, &path)?)
243                    } else {
244                        let file = system_path_to_file(db, &path).ok()?;
245                        (ModuleKind::Module, file)
246                    };
247                    Some(Module::file_module(
248                        db,
249                        file,
250                        resolver_environment,
251                        Cow::Owned(name),
252                        kind,
253                        module.search_path(db).clone(),
254                    ))
255                })
256                .collect()
257        }
258        SystemOrVendoredPathRef::Vendored(parent_directory) => db
259            .vendored()
260            .read_directory(parent_directory)
261            .filter(|entry| {
262                let ty = entry.file_type();
263                let path = entry.path();
264                is_submodule(
265                    ty.is_directory(),
266                    ty.is_file(),
267                    path.file_name(),
268                    path.extension(),
269                )
270            })
271            .filter_map(|entry| {
272                let stem = entry.path().file_stem()?;
273                let mut name = module.name(db).clone();
274                name.extend(&ModuleName::new(stem)?);
275
276                let (kind, file) = if entry.file_type().is_directory() {
277                    (
278                        ModuleKind::Package,
279                        find_package_init_vendored(db, entry.path())?,
280                    )
281                } else {
282                    let file = vendored_path_to_file(db, entry.path()).ok()?;
283                    (ModuleKind::Module, file)
284                };
285                Some(Module::file_module(
286                    db,
287                    file,
288                    resolver_environment,
289                    Cow::Owned(name),
290                    kind,
291                    module.search_path(db).clone(),
292                ))
293            })
294            .collect(),
295    })
296}
297
298/// A module that resolves to a file (`lib.py` or `package/__init__.py`).
299#[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)]
300pub struct FileModule<'db> {
301    #[returns(ref)]
302    pub(super) name: ModuleName,
303    #[returns(copy)]
304    pub(super) kind: ModuleKind,
305    #[returns(ref)]
306    pub(super) search_path: SearchPath,
307    #[returns(copy)]
308    pub(super) file: File,
309    #[returns(copy)]
310    pub(super) resolver_environment: ResolverEnvironment<'db>,
311    #[returns(copy)]
312    pub(super) known: Option<KnownModule>,
313}
314
315/// A namespace package.
316///
317/// Namespace packages are special because there are
318/// multiple possible paths and they have no corresponding code file.
319#[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)]
320pub struct NamespacePackage<'db> {
321    #[returns(copy)]
322    pub(super) resolver_environment: ResolverEnvironment<'db>,
323    #[returns(ref)]
324    pub(super) name: ModuleName,
325}
326
327#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, get_size2::GetSize)]
328pub enum ModuleKind {
329    /// A single-file module (e.g. `foo.py` or `foo.pyi`)
330    Module,
331
332    /// A python package (`foo/__init__.py` or `foo/__init__.pyi`)
333    Package,
334}
335
336impl ModuleKind {
337    pub const fn is_package(self) -> bool {
338        matches!(self, ModuleKind::Package)
339    }
340    pub const fn is_module(self) -> bool {
341        matches!(self, ModuleKind::Module)
342    }
343}
344
345/// Enumeration of modules in which types with dedicated semantic behavior are located.
346#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, strum_macros::EnumString, get_size2::GetSize)]
347#[cfg_attr(test, derive(strum_macros::EnumIter))]
348#[strum(serialize_all = "snake_case")]
349pub enum KnownModule {
350    Builtins,
351    Enum,
352    Types,
353    #[strum(serialize = "_typeshed")]
354    Typeshed,
355    TypingExtensions,
356    Typing,
357    Sys,
358    Os,
359    Tempfile,
360    Pathlib,
361    Datetime,
362    Decimal,
363    Ipaddress,
364    Re,
365    Abc,
366    Dataclasses,
367    Functools,
368    Collections,
369    #[strum(serialize = "collections.abc")]
370    CollectionsAbc,
371    #[strum(serialize = "_collections_abc")]
372    CollectionsAbcInternal,
373    Inspect,
374    #[strum(serialize = "string.templatelib")]
375    Templatelib,
376    #[strum(serialize = "_typeshed._type_checker_internals")]
377    TypeCheckerInternals,
378    TyExtensions,
379    #[strum(serialize = "ty_extensions._internal")]
380    TyExtensionsInternal,
381    #[strum(serialize = "ty_extensions.pydantic")]
382    TyExtensionsPydantic,
383    #[strum(serialize = "importlib")]
384    ImportLib,
385    #[strum(serialize = "unittest.mock")]
386    UnittestMock,
387    Uuid,
388    Warnings,
389    Numbers,
390    #[strum(serialize = "struct", serialize = "_struct")]
391    Struct,
392    // Third-party modules
393    #[strum(serialize = "pydantic.config")]
394    PydanticConfig,
395    #[strum(serialize = "pydantic.fields")]
396    PydanticFields,
397    #[strum(serialize = "pydantic.functional_validators")]
398    PydanticFunctionalValidators,
399    #[strum(serialize = "pydantic.main")]
400    PydanticMain,
401    #[strum(serialize = "pydantic.root_model")]
402    PydanticRootModel,
403    #[strum(serialize = "pydantic_settings.main")]
404    PydanticSettingsMain,
405    #[strum(serialize = "pydantic.types")]
406    PydanticTypes,
407    Pytest,
408    #[strum(serialize = "_pytest.fixtures")]
409    PytestFixtures,
410    #[strum(serialize = "_pytest.mark.structures")]
411    PytestMarkStructures,
412}
413
414impl KnownModule {
415    pub const fn as_str(self) -> &'static str {
416        match self {
417            Self::Builtins => "builtins",
418            Self::Enum => "enum",
419            Self::Types => "types",
420            Self::Typing => "typing",
421            Self::Typeshed => "_typeshed",
422            Self::TypingExtensions => "typing_extensions",
423            Self::Sys => "sys",
424            Self::Os => "os",
425            Self::Tempfile => "tempfile",
426            Self::Pathlib => "pathlib",
427            Self::Datetime => "datetime",
428            Self::Decimal => "decimal",
429            Self::Ipaddress => "ipaddress",
430            Self::Re => "re",
431            Self::Abc => "abc",
432            Self::Dataclasses => "dataclasses",
433            Self::Functools => "functools",
434            Self::Collections => "collections",
435            Self::CollectionsAbc => "collections.abc",
436            Self::CollectionsAbcInternal => "_collections_abc",
437            Self::Inspect => "inspect",
438            Self::TypeCheckerInternals => "_typeshed._type_checker_internals",
439            Self::TyExtensions => "ty_extensions",
440            Self::TyExtensionsInternal => "ty_extensions._internal",
441            Self::TyExtensionsPydantic => "ty_extensions.pydantic",
442            Self::ImportLib => "importlib",
443            Self::Warnings => "warnings",
444            Self::UnittestMock => "unittest.mock",
445            Self::Uuid => "uuid",
446            Self::Templatelib => "string.templatelib",
447            Self::Numbers => "numbers",
448            Self::Struct => "struct",
449            Self::PydanticConfig => "pydantic.config",
450            Self::PydanticFields => "pydantic.fields",
451            Self::PydanticFunctionalValidators => "pydantic.functional_validators",
452            Self::PydanticMain => "pydantic.main",
453            Self::PydanticRootModel => "pydantic.root_model",
454            Self::PydanticSettingsMain => "pydantic_settings.main",
455            Self::PydanticTypes => "pydantic.types",
456            Self::Pytest => "pytest",
457            Self::PytestFixtures => "_pytest.fixtures",
458            Self::PytestMarkStructures => "_pytest.mark.structures",
459        }
460    }
461
462    pub fn name(self) -> ModuleName {
463        ModuleName::new_static(self.as_str())
464            .unwrap_or_else(|| panic!("{self} should be a valid module name!"))
465    }
466
467    fn try_from_search_path_and_name(search_path: &SearchPath, name: &ModuleName) -> Option<Self> {
468        let known_module = Self::from_str(name.as_str()).ok()?;
469
470        let is_expected_search_path = if known_module.is_third_party() {
471            search_path.can_contain_third_party_code()
472        } else {
473            search_path.is_standard_library()
474        };
475
476        is_expected_search_path.then_some(known_module)
477    }
478
479    /// Return `true` if this module is provided by a supported third-party package.
480    pub const fn is_third_party(self) -> bool {
481        match self {
482            Self::PydanticConfig
483            | Self::PydanticFields
484            | Self::PydanticFunctionalValidators
485            | Self::PydanticMain
486            | Self::PydanticRootModel
487            | Self::PydanticSettingsMain
488            | Self::PydanticTypes
489            | Self::Pytest
490            | Self::PytestFixtures
491            | Self::PytestMarkStructures => true,
492            Self::Builtins
493            | Self::Enum
494            | Self::Types
495            | Self::Typeshed
496            | Self::TypingExtensions
497            | Self::Typing
498            | Self::Sys
499            | Self::Os
500            | Self::Tempfile
501            | Self::Pathlib
502            | Self::Datetime
503            | Self::Decimal
504            | Self::Ipaddress
505            | Self::Re
506            | Self::Abc
507            | Self::Dataclasses
508            | Self::Functools
509            | Self::Collections
510            | Self::CollectionsAbc
511            | Self::CollectionsAbcInternal
512            | Self::Inspect
513            | Self::Templatelib
514            | Self::TypeCheckerInternals
515            | Self::TyExtensions
516            | Self::TyExtensionsInternal
517            | Self::TyExtensionsPydantic
518            | Self::ImportLib
519            | Self::UnittestMock
520            | Self::Uuid
521            | Self::Warnings
522            | Self::Numbers
523            | Self::Struct => false,
524        }
525    }
526
527    pub const fn is_builtins(self) -> bool {
528        matches!(self, Self::Builtins)
529    }
530
531    pub const fn is_typing(self) -> bool {
532        matches!(self, Self::Typing)
533    }
534
535    pub const fn is_typing_extensions(self) -> bool {
536        matches!(self, Self::TypingExtensions)
537    }
538
539    pub const fn is_ty_extensions(self) -> bool {
540        matches!(self, Self::TyExtensions)
541    }
542
543    pub const fn is_ty_extensions_internal(self) -> bool {
544        matches!(self, Self::TyExtensionsInternal)
545    }
546
547    pub const fn is_inspect(self) -> bool {
548        matches!(self, Self::Inspect)
549    }
550
551    pub const fn is_importlib(self) -> bool {
552        matches!(self, Self::ImportLib)
553    }
554
555    pub const fn is_functools(self) -> bool {
556        matches!(self, Self::Functools)
557    }
558
559    pub const fn is_dataclasses(self) -> bool {
560        matches!(self, Self::Dataclasses)
561    }
562}
563
564impl std::fmt::Display for KnownModule {
565    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
566        f.write_str(self.as_str())
567    }
568}
569
570#[cfg(test)]
571mod tests {
572    use super::*;
573    use strum::IntoEnumIterator;
574
575    #[test]
576    fn known_module_roundtrip_from_str() {
577        let stdlib_search_path = SearchPath::vendored_stdlib();
578
579        for module in KnownModule::iter().filter(|module| !module.is_third_party()) {
580            let module_name = module.name();
581
582            assert_eq!(
583                KnownModule::try_from_search_path_and_name(&stdlib_search_path, &module_name),
584                Some(module),
585                "The strum `EnumString` implementation appears to be incorrect for `{module_name}`"
586            );
587        }
588    }
589}
ty_module_resolver/module.rs

ty_module_resolver/
module.rs
