Harmonise methods for distinguishing different Python source types

Author: AlexWaygood

Cargo.lock

@@ -3,6 +3,7 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
 use lsp_types::Url;
+use ruff_python_ast::PySourceType;
 use rustc_hash::FxHashMap;
 
 use crate::{
@@ -72,10 +73,7 @@ impl Index {
     pub(crate) fn key_from_url(&self, url: Url) -> DocumentKey {
         if self.notebook_cells.contains_key(&url) {
             DocumentKey::NotebookCell(url)
-        } else if Path::new(url.path())
-            .extension()
-            .map_or(false, |ext| ext.eq_ignore_ascii_case("ipynb"))
-        {
+        } else if PySourceType::from(url.path()).is_ipynb() {
             DocumentKey::Notebook(url)
         } else {
             DocumentKey::Text(url)

crates/red_knot_server/src/session/index.rs

@@ -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/Cargo.toml

@@ -7,7 +7,7 @@ use std::path::PathBuf;
 use anyhow::Result;
 
 use ruff_python_codegen::round_trip;
-use ruff_python_stdlib::path::is_jupyter_notebook;
+use ruff_python_ast::PySourceType;
 
 #[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_dev/src/round_trip.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,10 @@ pub(crate) fn builtin_module_shadowing(
     allowed_modules: &[String],
     target_version: PythonVersion,
 ) -> Option<Diagnostic> {
-    if !path
-        .extension()
-        .is_some_and(|ext| ext == "py" || ext == "pyi")
-    {
+    if !matches!(
+        PySourceType::try_from_path(path),
+        Some(PySourceType::Python | PySourceType::Stub)
+    ) {
         return None;
     }
 

crates/ruff_linter/src/rules/flake8_builtins/rules/builtin_module_shadowing.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(|source_type|source_type.is_python())
         // Ignore any files that are direct children of the project root.
         && !path
             .parent()

crates/ruff_linter/src/rules/flake8_no_pep420/rules/implicit_namespace_package.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,10 @@ pub(crate) fn invalid_module_name(
     package: Option<&Path>,
     ignore_names: &IgnoreNames,
 ) -> Option<Diagnostic> {
-    if !path
-        .extension()
-        .is_some_and(|ext| ext == "py" || ext == "pyi")
-    {
+    if !matches!(
+        PySourceType::try_from_path(path),
+        Some(PySourceType::Python | PySourceType::Stub)
+    ) {
         return None;
     }
 

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

@@ -99,13 +99,17 @@ impl PySourceType {
 
         Some(ty)
     }
-}
 
-impl<P: AsRef<Path>> From<P> for PySourceType {
-    fn from(path: P) -> Self {
+    pub fn try_from_path(path: impl AsRef<Path>) -> Option<Self> {
         path.as_ref()
             .extension()
             .and_then(OsStr::to_str)
-            .map_or(Self::Python, Self::from_extension)
+            .and_then(Self::try_from_extension)
+    }
+}
+
+impl<P: AsRef<Path>> From<P> for PySourceType {
+    fn from(path: P) -> Self {
+        Self::try_from_path(path).unwrap_or_default()
     }
 }

crates/ruff_python_ast/src/lib.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_stdlib::path::is_python_stub_file;
+use ruff_python_ast::{self as ast, Expr, ExprContext, Operator, PySourceType, Stmt};
 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_semantic/src/model.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| {
-        file_name == "__init__.py"
-            || file_name == "__init__.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));
-    }
+    matches!(
+        path.file_name().and_then(OsStr::to_str),
+        Some("__init__.py" | "__init__.pyi" | "__main__.py" | "__main__.pyi")
+    )
 }