Make SourceKind a required parameter

crates/ruff/src/checkers/imports.rs

@@ -85,7 +85,7 @@ pub(crate) fn check_imports(
  stylist: &Stylist,
  path: &Path,
  package: Option<&Path>,
- source_kind: Option<&SourceKind>,
+ source_kind: &SourceKind,
  source_type: PySourceType,
 ) -> (Vec<Diagnostic>, Option<ImportMap>) {
  // Extract all import blocks from the AST.

crates/ruff/src/linter.rs

@@ -81,7 +81,7 @@ pub fn check_path(
  directives: &Directives,
  settings: &Settings,
  noqa: flags::Noqa,
- source_kind: Option<&SourceKind>,
+ source_kind: &SourceKind,
  source_type: PySourceType,
 ) -> LinterResult<(Vec<Diagnostic>, Option<ImportMap>)> {
  // Aggregate all diagnostics.
@@ -270,17 +270,17 @@ const MAX_ITERATIONS: usize = 100;
 pub fn add_noqa_to_path(
  path: &Path,
  package: Option<&Path>,
+ source_kind: &SourceKind,
  source_type: PySourceType,
  settings: &Settings,
 ) -> Result<usize> {
- // Read the file from disk.
- let contents = std::fs::read_to_string(path)?;
+ let contents = source_kind.source_code();
 
  // Tokenize once.
- let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, source_type.as_mode());
+ let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, source_type.as_mode());
 
  // Map row and column locations to byte slices (lazily).
- let locator = Locator::new(&contents);
+ let locator = Locator::new(contents);
 
  // Detect the current code style (lazily).
  let stylist = Stylist::from_tokens(&tokens, &locator);
@@ -310,21 +310,20 @@ pub fn add_noqa_to_path(
  &directives,
  settings,
  flags::Noqa::Disabled,
- None,
+ source_kind,
  source_type,
  );
 
  // Log any parse errors.
  if let Some(err) = error {
- // TODO(dhruvmanila): This should use `SourceKind`, update when
- // `--add-noqa` is supported for Jupyter notebooks.
  error!(
  "{}",
- DisplayParseError::new(err, locator.to_source_code(), None)
+ DisplayParseError::new(err, locator.to_source_code(), source_kind)
  );
  }
 
  // Add any missing `# noqa` pragmas.
+ // TODO(dhruvmanila): Add support for Jupyter Notebooks
  add_noqa(
  path,
  &diagnostics.0,
@@ -377,7 +376,7 @@ pub fn lint_only(
  &directives,
  settings,
  noqa,
- Some(source_kind),
+ source_kind,
  source_type,
  );
 
@@ -471,7 +470,7 @@ pub fn lint_fix<'a>(
  &directives,
  settings,
  noqa,
- Some(source_kind),
+ source_kind,
  source_type,
  );
 

crates/ruff/src/logging.rs

@@ -139,14 +139,14 @@ pub fn set_up_logging(level: &LogLevel) -> Result<()> {
 pub struct DisplayParseError<'a> {
  error: ParseError,
  source_code: SourceCode<'a, 'a>,
- source_kind: Option<&'a SourceKind>,
+ source_kind: &'a SourceKind,
 }
 
 impl<'a> DisplayParseError<'a> {
  pub fn new(
  error: ParseError,
  source_code: SourceCode<'a, 'a>,
- source_kind: Option<&'a SourceKind>,
+ source_kind: &'a SourceKind,
  ) -> Self {
  Self {
  error,
@@ -171,32 +171,29 @@ impl Display for DisplayParseError<'_> {
  // If we're working on a Jupyter notebook, translate the positions
  // with respect to the cell and row in the cell. This is the same
  // format as the `TextEmitter`.
- let error_location = if let Some(jupyter_index) = self
- .source_kind
- .and_then(SourceKind::notebook)
- .map(Notebook::index)
- {
- write!(
- f,
- "cell {cell}{colon}",
- cell = jupyter_index
- .cell(source_location.row.get())
- .unwrap_or_default(),
- colon = ":".cyan(),
- )?;
+ let error_location =
+ if let Some(jupyter_index) = self.source_kind.as_ipy_notebook().map(Notebook::index) {
+ write!(
+ f,
+ "cell {cell}{colon}",
+ cell = jupyter_index
+ .cell(source_location.row.get())
+ .unwrap_or_default(),
+ colon = ":".cyan(),
+ )?;
 
- SourceLocation {
- row: OneIndexed::new(
- jupyter_index
- .cell_row(source_location.row.get())
- .unwrap_or(1) as usize,
- )
- .unwrap(),
- column: source_location.column,
- }
- } else {
- source_location
- };
+  SourceLocation {
+  row: OneIndexed::new(
+  jupyter_index
+  .cell_row(source_location.row.get())
+  .unwrap_or(1) as usize,
+  )
+  .unwrap(),
+  column: source_location.column,
+  }
+  } else {
+  source_location
+  };
 
  write!(
  f,

crates/ruff/src/rules/isort/block.rs

@@ -43,7 +43,7 @@ impl<'a> BlockBuilder<'a> {
  locator: &'a Locator<'a>,
  directives: &'a IsortDirectives,
  is_stub: bool,
- source_kind: Option<&'a SourceKind>,
+ source_kind: &'a SourceKind,
  ) -> Self {
  Self {
  locator,
@@ -53,7 +53,7 @@ impl<'a> BlockBuilder<'a> {
  exclusions: &directives.exclusions,
  nested: false,
  cell_offsets: source_kind
- .and_then(SourceKind::notebook)
+ .as_ipy_notebook()
  .map(Notebook::cell_offsets)
  .map(|offsets| offsets.iter().peekable()),
  }

crates/ruff/src/rules/pyflakes/mod.rs

@@ -27,6 +27,7 @@ mod tests {
  use crate::registry::{AsRule, Linter, Rule};
  use crate::rules::pyflakes;
  use crate::settings::{flags, Settings};
+ use crate::source_kind::SourceKind;
  use crate::test::{test_path, test_snippet};
  use crate::{assert_messages, directives};
 
@@ -508,6 +509,7 @@ mod tests {
  fn flakes(contents: &str, expected: &[Rule]) {
  let contents = dedent(contents);
  let source_type = PySourceType::default();
+ let source_kind = SourceKind::Python(contents.to_string());
  let settings = Settings::for_rules(Linter::Pyflakes.rules());
  let tokens: Vec<LexResult> = ruff_python_parser::tokenize(&contents, source_type.as_mode());
  let locator = Locator::new(&contents);
@@ -532,7 +534,7 @@ mod tests {
  &directives,
  &settings,
  flags::Noqa::Enabled,
- None,
+ &source_kind,
  source_type,
  );
  diagnostics.sort_by_key(Ranged::start);

crates/ruff/src/source_kind.rs

@@ -10,15 +10,6 @@ pub enum SourceKind {
 }
 
 impl SourceKind {
- /// Return the [`Notebook`] if the source kind is [`SourceKind::IpyNotebook`].
- pub fn notebook(&self) -> Option<&Notebook> {
- if let Self::IpyNotebook(notebook) = self {
- Some(notebook)
- } else {
- None
- }
- }
-
  #[must_use]
  pub(crate) fn updated(&self, new_source: String, source_map: &SourceMap) -> Self {
  match self {

crates/ruff/src/test.rs

@@ -132,7 +132,7 @@ pub(crate) fn test_contents<'a>(
  &directives,
  settings,
  flags::Noqa::Enabled,
- Some(source_kind),
+ source_kind,
  source_type,
  );
 
@@ -195,7 +195,7 @@ pub(crate) fn test_contents<'a>(
  &directives,
  settings,
  flags::Noqa::Enabled,
- Some(source_kind),
+ source_kind,
  source_type,
  );
 
@@ -274,7 +274,7 @@ fn print_diagnostics(diagnostics: Vec<Diagnostic>, path: &Path, source: &SourceK
  })
  .collect();
 
- if let Some(notebook) = source.notebook() {
+ if let Some(notebook) = source.as_ipy_notebook() {
  print_jupyter_messages(&messages, path, notebook)
  } else {
  print_messages(&messages)

crates/ruff_cli/src/commands/add_noqa.rs

@@ -12,6 +12,7 @@ use ruff_python_ast::{PySourceType, SourceType};
 use ruff_workspace::resolver::{python_files_in_path, PyprojectConfig};
 
 use crate::args::Overrides;
+use crate::diagnostics::LintSource;
 
 /// Add `noqa` directives to a collection of files.
 pub(crate) fn add_noqa(
@@ -56,7 +57,15 @@ pub(crate) fn add_noqa(
  .and_then(|parent| package_roots.get(parent))
  .and_then(|package| *package);
  let settings = resolver.resolve(path, pyproject_config);
- match add_noqa_to_path(path, package, source_type, settings) {
+ let LintSource(source_kind) = match LintSource::try_from_path(path, source_type) {
+ Ok(Some(source)) => source,
+ Ok(None) => return None,
+ Err(e) => {
+ error!("Failed to extract source from {}: {e}", path.display());
+ return None;
+ }
+ };
+ match add_noqa_to_path(path, package, &source_kind, source_type, settings) {
  Ok(count) => Some(count),
  Err(e) => {
  error!("Failed to add noqa to {}: {e}", path.display());

crates/ruff_cli/src/diagnostics.rs

@@ -354,7 +354,7 @@ pub(crate) fn lint_path(
  source_kind.source_code(),
  &LineIndex::from_source_text(source_kind.source_code())
  ),
- Some(&source_kind),
+ &source_kind,
  )
  );
  }
@@ -503,11 +503,11 @@ pub(crate) fn lint_stdin(
 }
 
 #[derive(Debug)]
-struct LintSource(SourceKind);
+pub(crate) struct LintSource(pub(crate) SourceKind);
 
 impl LintSource {
  /// Extract the lint [`LintSource`] from the given file path.
- fn try_from_path(
+ pub(crate) fn try_from_path(
  path: &Path,
  source_type: PySourceType,
  ) -> Result<Option<LintSource>, SourceExtractionError> {
@@ -526,7 +526,7 @@ impl LintSource {
  /// Extract the lint [`LintSource`] from the raw string contents, optionally accompanied by a
  /// file path indicating the path to the file from which the contents were read. If provided,
  /// the file path should be used for diagnostics, but not for reading the file from disk.
- fn try_from_source_code(
+ pub(crate) fn try_from_source_code(
  source_code: String,
  source_type: PySourceType,
  ) -> Result<Option<LintSource>, SourceExtractionError> {

crates/ruff_wasm/src/lib.rs

@@ -10,6 +10,7 @@ use ruff::linter::{check_path, LinterResult};
 use ruff::registry::AsRule;
 use ruff::settings::types::PythonVersion;
 use ruff::settings::{defaults, flags, Settings};
+use ruff::source_kind::SourceKind;
 use ruff_formatter::{FormatResult, Formatted};
 use ruff_python_ast::{Mod, PySourceType};
 use ruff_python_codegen::Stylist;
@@ -165,6 +166,9 @@ impl Workspace {
  pub fn check(&self, contents: &str) -> Result<JsValue, Error> {
  let source_type = PySourceType::default();
 
+ // TODO(dhruvmanila): Support Jupyter Notebooks
+ let source_kind = SourceKind::Python(contents.to_string());
+
  // Tokenize once.
  let tokens: Vec<LexResult> = ruff_python_parser::tokenize(contents, source_type.as_mode());
 
@@ -195,7 +199,7 @@ impl Workspace {
  &directives,
  &self.settings,
  flags::Noqa::Enabled,
- None,
+ &source_kind,
  source_type,
  );