Harmonise methods for distinguishing different Python source types (#13682)

Cargo.lock

@@ -2466,7 +2466,6 @@ dependencies = [
  "ruff_python_codegen",
  "ruff_python_formatter",
  "ruff_python_parser",
- "ruff_python_stdlib",
  "ruff_python_trivia",
  "ruff_workspace",
  "schemars",

crates/ruff_dev/Cargo.toml

@@ -20,7 +20,6 @@ ruff_python_ast = { workspace = true }
 ruff_python_codegen = { workspace = true }
 ruff_python_formatter = { workspace = true }
 ruff_python_parser = { workspace = true }
-ruff_python_stdlib = { workspace = true }
 ruff_python_trivia = { workspace = true }
 ruff_workspace = { workspace = true, features = ["schemars"] }
 

crates/ruff_dev/src/round_trip.rs

@@ -6,8 +6,8 @@ use std::path::PathBuf;
 
 use anyhow::Result;
 
+use ruff_python_ast::PySourceType;
 use ruff_python_codegen::round_trip;
-use ruff_python_stdlib::path::is_jupyter_notebook;
 
 #[derive(clap::Args)]
 pub(crate) struct Args {
@@ -18,7 +18,7 @@ pub(crate) struct Args {
 
 pub(crate) fn main(args: &Args) -> Result<()> {
     let path = args.file.as_path();
-    if is_jupyter_notebook(path) {
+    if PySourceType::from(path).is_ipynb() {
         println!("{}", ruff_notebook::round_trip(path)?);
     } else {
         let contents = fs::read_to_string(&args.file)?;

crates/ruff_linter/src/rules/flake8_builtins/rules/builtin_module_shadowing.rs

@@ -2,6 +2,7 @@ use std::path::Path;
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::PySourceType;
 use ruff_python_stdlib::path::is_module_file;
 use ruff_python_stdlib::sys::is_known_standard_library;
 use ruff_text_size::TextRange;
@@ -42,10 +43,7 @@ pub(crate) fn builtin_module_shadowing(
     allowed_modules: &[String],
     target_version: PythonVersion,
 ) -> Option<Diagnostic> {
-    if !path
+    if !PySourceType::try_from_path(path).is_some_and(PySourceType::is_py_file_or_stub) {
-        .extension()
-        .is_some_and(|ext| ext == "py" || ext == "pyi")
-    {
         return None;
     }
 

crates/ruff_linter/src/rules/flake8_no_pep420/rules/implicit_namespace_package.rs

@@ -2,6 +2,7 @@ use std::path::{Path, PathBuf};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::PySourceType;
 use ruff_python_trivia::CommentRanges;
 use ruff_source_file::Locator;
 use ruff_text_size::{TextRange, TextSize};
@@ -51,7 +52,7 @@ pub(crate) fn implicit_namespace_package(
 ) -> Option<Diagnostic> {
     if package.is_none()
         // Ignore non-`.py` files, which don't require an `__init__.py`.
-        && path.extension().is_some_and( |ext| ext == "py")
+        && PySourceType::try_from_path(path).is_some_and(PySourceType::is_py_file)
         // Ignore any files that are direct children of the project root.
         && !path
             .parent()

crates/ruff_linter/src/rules/pep8_naming/rules/invalid_module_name.rs

@@ -3,6 +3,7 @@ use std::path::Path;
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::PySourceType;
 use ruff_python_stdlib::identifiers::{is_migration_name, is_module_name};
 use ruff_python_stdlib::path::is_module_file;
 use ruff_text_size::TextRange;
@@ -53,10 +54,7 @@ pub(crate) fn invalid_module_name(
     package: Option<&Path>,
     ignore_names: &IgnoreNames,
 ) -> Option<Diagnostic> {
-    if !path
+    if !PySourceType::try_from_path(path).is_some_and(PySourceType::is_py_file_or_stub) {
-        .extension()
-        .is_some_and(|ext| ext == "py" || ext == "pyi")
-    {
         return None;
     }
 

crates/ruff_python_ast/src/lib.rs

@@ -68,7 +68,7 @@ pub enum TomlSourceType {
     Unrecognized,
 }
 
-#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, is_macro::Is)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub enum PySourceType {
     /// The source is a Python file (`.py`).
@@ -99,13 +99,33 @@ impl PySourceType {
 
         Some(ty)
     }
+
+    pub fn try_from_path(path: impl AsRef<Path>) -> Option<Self> {
+        path.as_ref()
+            .extension()
+            .and_then(OsStr::to_str)
+            .and_then(Self::try_from_extension)
+    }
+
+    pub const fn is_py_file(self) -> bool {
+        matches!(self, Self::Python)
+    }
+
+    pub const fn is_stub(self) -> bool {
+        matches!(self, Self::Stub)
+    }
+
+    pub const fn is_py_file_or_stub(self) -> bool {
+        matches!(self, Self::Python | Self::Stub)
+    }
+
+    pub const fn is_ipynb(self) -> bool {
+        matches!(self, Self::Ipynb)
+    }
 }
 
 impl<P: AsRef<Path>> From<P> for PySourceType {
     fn from(path: P) -> Self {
-        path.as_ref()
+        Self::try_from_path(path).unwrap_or_default()
-            .extension()
-            .and_then(OsStr::to_str)
-            .map_or(Self::Python, Self::from_extension)
     }
 }

crates/ruff_python_semantic/src/model.rs

@@ -5,8 +5,7 @@ use rustc_hash::FxHashMap;
 
 use ruff_python_ast::helpers::from_relative_import;
 use ruff_python_ast::name::{QualifiedName, UnqualifiedName};
-use ruff_python_ast::{self as ast, Expr, ExprContext, Operator, Stmt};
+use ruff_python_ast::{self as ast, Expr, ExprContext, Operator, PySourceType, Stmt};
-use ruff_python_stdlib::path::is_python_stub_file;
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::binding::{
@@ -2246,7 +2245,7 @@ bitflags! {
 
 impl SemanticModelFlags {
     pub fn new(path: &Path) -> Self {
-        if is_python_stub_file(path) {
+        if PySourceType::from(path).is_stub() {
             Self::STUB_FILE
         } else {
             Self::default()

crates/ruff_python_stdlib/src/path.rs

@@ -1,3 +1,4 @@
+use std::ffi::OsStr;
 use std::path::Path;
 
 /// Return `true` if the [`Path`] is named `pyproject.toml`.
@@ -6,38 +7,10 @@ pub fn is_pyproject_toml(path: &Path) -> bool {
         .is_some_and(|name| name == "pyproject.toml")
 }
 
-/// Return `true` if the [`Path`] appears to be that of a Python interface definition file (`.pyi`).
-pub fn is_python_stub_file(path: &Path) -> bool {
-    path.extension().is_some_and(|ext| ext == "pyi")
-}
-
-/// Return `true` if the [`Path`] appears to be that of a Jupyter notebook (`.ipynb`).
-pub fn is_jupyter_notebook(path: &Path) -> bool {
-    path.extension().is_some_and(|ext| ext == "ipynb")
-}
-
 /// Return `true` if a [`Path`] should use the name of its parent directory as its module name.
 pub fn is_module_file(path: &Path) -> bool {
-    path.file_name().is_some_and(|file_name| {
+    matches!(
-        file_name == "__init__.py"
+        path.file_name().and_then(OsStr::to_str),
-            || file_name == "__init__.pyi"
+        Some("__init__.py" | "__init__.pyi" | "__main__.py" | "__main__.pyi")
-            || file_name == "__main__.py"
+    )
-            || file_name == "__main__.pyi"
-    })
-}
-
-#[cfg(test)]
-mod tests {
-    use std::path::Path;
-
-    use crate::path::is_jupyter_notebook;
-
-    #[test]
-    fn test_is_jupyter_notebook() {
-        let path = Path::new("foo/bar/baz.ipynb");
-        assert!(is_jupyter_notebook(path));
-
-        let path = Path::new("foo/bar/baz.py");
-        assert!(!is_jupyter_notebook(path));
-    }
 }