Skip to main content

ruff_python_ast/
lib.rs

1use std::ffi::OsStr;
2use std::path::Path;
3
4pub use expression::*;
5pub use generated::*;
6pub use int::*;
7pub use node_index::*;
8pub use nodes::*;
9pub use operator_precedence::*;
10pub use python_version::*;
11
12pub mod comparable;
13pub mod docstrings;
14mod expression;
15pub mod find_node;
16mod generated;
17pub mod helpers;
18pub mod identifier;
19mod int;
20pub mod name;
21mod node;
22mod node_index;
23mod nodes;
24pub mod operator_precedence;
25pub mod parenthesize;
26mod python_version;
27pub mod relocate;
28pub mod script;
29pub mod statement_visitor;
30pub mod stmt_if;
31pub mod str;
32pub mod str_prefix;
33pub mod token;
34pub mod traversal;
35pub mod types;
36pub mod visitor;
37pub mod whitespace;
38
39/// The type of a source file.
40#[derive(Clone, Copy, Debug, PartialEq, is_macro::Is)]
41pub enum SourceType {
42    /// The file contains Python source code.
43    Python(PySourceType),
44    /// The file contains TOML.
45    Toml(TomlSourceType),
46    /// The file contains Markdown.
47    Markdown,
48}
49
50impl Default for SourceType {
51    fn default() -> Self {
52        Self::Python(PySourceType::Python)
53    }
54}
55
56impl<P: AsRef<Path>> From<P> for SourceType {
57    fn from(path: P) -> Self {
58        match path.as_ref().file_name() {
59            Some(filename) if filename == "pyproject.toml" => Self::Toml(TomlSourceType::Pyproject),
60            Some(filename) if filename == "Pipfile" => Self::Toml(TomlSourceType::Pipfile),
61            Some(filename) if filename == "poetry.lock" => Self::Toml(TomlSourceType::Poetry),
62            _ => match path.as_ref().extension() {
63                Some(ext) if ext == "toml" => Self::Toml(TomlSourceType::Unrecognized),
64                Some(ext) if ext == "md" || ext == "qmd" => Self::Markdown,
65                _ => Self::Python(PySourceType::from(path)),
66            },
67        }
68    }
69}
70
71#[derive(Clone, Copy, Debug, PartialEq, is_macro::Is)]
72pub enum TomlSourceType {
73    /// The source is a `pyproject.toml`.
74    Pyproject,
75    /// The source is a `Pipfile`.
76    Pipfile,
77    /// The source is a `poetry.lock`.
78    Poetry,
79    /// The source is an unrecognized TOML file.
80    Unrecognized,
81}
82
83#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
84#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
85pub enum PySourceType {
86    /// The source is a Python file (`.py`, `.pyw`).
87    /// Note: `.pyw` files contain Python code, but do not represent importable namespaces.
88    /// Consider adding a separate source type later if combining the two causes issues.
89    #[default]
90    Python,
91    /// The source is a Python stub file (`.pyi`).
92    Stub,
93    /// The source is a Jupyter notebook (`.ipynb`).
94    Ipynb,
95}
96
97impl PySourceType {
98    /// Infers the source type from the file extension.
99    ///
100    /// Falls back to `Python` if the extension is not recognized.
101    pub fn from_extension(extension: &str) -> Self {
102        Self::try_from_extension(extension).unwrap_or_default()
103    }
104
105    /// Infers the source type from the file extension.
106    pub fn try_from_extension(extension: &str) -> Option<Self> {
107        let ty = match extension {
108            "py" => Self::Python,
109            "pyi" => Self::Stub,
110            "pyw" => Self::Python,
111            "ipynb" => Self::Ipynb,
112            _ => return None,
113        };
114
115        Some(ty)
116    }
117
118    pub fn try_from_path(path: impl AsRef<Path>) -> Option<Self> {
119        path.as_ref()
120            .extension()
121            .and_then(OsStr::to_str)
122            .and_then(Self::try_from_extension)
123    }
124
125    pub const fn is_py_file(self) -> bool {
126        matches!(self, Self::Python)
127    }
128
129    pub const fn is_stub(self) -> bool {
130        matches!(self, Self::Stub)
131    }
132
133    pub const fn is_py_file_or_stub(self) -> bool {
134        matches!(self, Self::Python | Self::Stub)
135    }
136
137    pub const fn is_ipynb(self) -> bool {
138        matches!(self, Self::Ipynb)
139    }
140}
141
142impl<P: AsRef<Path>> From<P> for PySourceType {
143    fn from(path: P) -> Self {
144        Self::try_from_path(path).unwrap_or_default()
145    }
146}