Skip to main content

rustpython_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 SourceType {
51    pub fn from_extension(ext: &str) -> Self {
52        match ext {
53            "toml" => Self::Toml(TomlSourceType::Unrecognized),
54            "md" => Self::Markdown,
55            _ => Self::Python(PySourceType::from_extension(ext)),
56        }
57    }
58}
59
60impl Default for SourceType {
61    fn default() -> Self {
62        Self::Python(PySourceType::Python)
63    }
64}
65
66impl<P: AsRef<Path>> From<P> for SourceType {
67    fn from(path: P) -> Self {
68        match path.as_ref().file_name() {
69            Some(filename) if filename == "pyproject.toml" => Self::Toml(TomlSourceType::Pyproject),
70            Some(filename) if filename == "Pipfile" => Self::Toml(TomlSourceType::Pipfile),
71            Some(filename) if filename == "poetry.lock" => Self::Toml(TomlSourceType::Poetry),
72            _ => Self::from_extension(
73                path.as_ref()
74                    .extension()
75                    .and_then(OsStr::to_str)
76                    .unwrap_or(""),
77            ),
78        }
79    }
80}
81
82#[derive(Clone, Copy, Debug, PartialEq, is_macro::Is)]
83pub enum TomlSourceType {
84    /// The source is a `pyproject.toml`.
85    Pyproject,
86    /// The source is a `Pipfile`.
87    Pipfile,
88    /// The source is a `poetry.lock`.
89    Poetry,
90    /// The source is an unrecognized TOML file.
91    Unrecognized,
92}
93
94#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
95#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
96pub enum PySourceType {
97    /// The source is a Python file (`.py`, `.pyw`).
98    /// Note: `.pyw` files contain Python code, but do not represent importable namespaces.
99    /// Consider adding a separate source type later if combining the two causes issues.
100    #[default]
101    Python,
102    /// The source is a Python stub file (`.pyi`).
103    Stub,
104    /// The source is a Jupyter notebook (`.ipynb`).
105    Ipynb,
106}
107
108impl PySourceType {
109    /// Infers the source type from the file extension.
110    ///
111    /// Falls back to `Python` if the extension is not recognized.
112    pub fn from_extension(extension: &str) -> Self {
113        Self::try_from_extension(extension).unwrap_or_default()
114    }
115
116    /// Infers the source type from the file extension.
117    pub fn try_from_extension(extension: &str) -> Option<Self> {
118        let ty = match extension {
119            "py" => Self::Python,
120            "pyi" => Self::Stub,
121            "pyw" => Self::Python,
122            "ipynb" => Self::Ipynb,
123            _ => return None,
124        };
125
126        Some(ty)
127    }
128
129    pub fn try_from_path(path: impl AsRef<Path>) -> Option<Self> {
130        path.as_ref()
131            .extension()
132            .and_then(OsStr::to_str)
133            .and_then(Self::try_from_extension)
134    }
135
136    pub const fn is_py_file(self) -> bool {
137        matches!(self, Self::Python)
138    }
139
140    pub const fn is_stub(self) -> bool {
141        matches!(self, Self::Stub)
142    }
143
144    pub const fn is_py_file_or_stub(self) -> bool {
145        matches!(self, Self::Python | Self::Stub)
146    }
147
148    pub const fn is_ipynb(self) -> bool {
149        matches!(self, Self::Ipynb)
150    }
151}
152
153impl<P: AsRef<Path>> From<P> for PySourceType {
154    fn from(path: P) -> Self {
155        Self::try_from_path(path).unwrap_or_default()
156    }
157}