Skip to content

Commit

Permalink
Prevent invalid utf8 indexing in cell magic detection (#9146)
Browse files Browse the repository at this point in the history
The example below used to panic because we tried to split at 2 bytes in
the 4-bytes character `转`.
```python
def sample_func(xx):
    """
    转置 (transpose)
    """
    return xx.T
```

Fixes #9145
Fixes astral-sh/ruff-vscode#362

The second commit is a small test refactoring.
  • Loading branch information
konstin authored Dec 15, 2023
1 parent 3ce145c commit cd3c2f7
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 54 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
{
"execution_count": null,
"cell_type": "code",
"id": "1",
"metadata": {},
"outputs": [],
"source": [
"def sample_func(xx):\n",
" \"\"\"\n",
" 转置 (transpose)\n",
" \"\"\"\n",
" return xx.T",
"# https://github.com/astral-sh/ruff-vscode/issues/362",
"DEFAULT_SYSTEM_PROMPT = (",
" \"Ты — Сайга, русскоязычный автоматический ассистент. \"",
" \"Ты разговариваешь с людьми и помогаешь им.\"",
")"
]
}
79 changes: 37 additions & 42 deletions crates/ruff_notebook/src/cell.rs
Original file line number Diff line number Diff line change
Expand Up @@ -171,48 +171,43 @@ impl Cell {

// Detect cell magics (which operate on multiple lines).
lines.any(|line| {
line.split_whitespace().next().is_some_and(|first| {
if first.len() < 2 {
return false;
}
let (token, command) = first.split_at(2);
// These cell magics are special in that the lines following them are valid
// Python code and the variables defined in that scope are available to the
// rest of the notebook.
//
// For example:
//
// Cell 1:
// ```python
// x = 1
// ```
//
// Cell 2:
// ```python
// %%time
// y = x
// ```
//
// Cell 3:
// ```python
// print(y) # Here, `y` is available.
// ```
//
// This is to avoid false positives when these variables are referenced
// elsewhere in the notebook.
token == "%%"
&& !matches!(
command,
"capture"
| "debug"
| "prun"
| "pypy"
| "python"
| "python3"
| "time"
| "timeit"
)
})
let Some(first) = line.split_whitespace().next() else {
return false;
};
if first.len() < 2 {
return false;
}
let Some(command) = first.strip_prefix("%%") else {
return false;
};
// These cell magics are special in that the lines following them are valid
// Python code and the variables defined in that scope are available to the
// rest of the notebook.
//
// For example:
//
// Cell 1:
// ```python
// x = 1
// ```
//
// Cell 2:
// ```python
// %%time
// y = x
// ```
//
// Cell 3:
// ```python
// print(y) # Here, `y` is available.
// ```
//
// This is to avoid false positives when these variables are referenced
// elsewhere in the notebook.
!matches!(
command,
"capture" | "debug" | "prun" | "pypy" | "python" | "python3" | "time" | "timeit"
)
})
}
}
Expand Down
28 changes: 16 additions & 12 deletions crates/ruff_notebook/src/notebook.rs
Original file line number Diff line number Diff line change
Expand Up @@ -421,25 +421,29 @@ mod tests {
));
}

#[test_case(Path::new("markdown.json"), false; "markdown")]
#[test_case(Path::new("only_magic.json"), true; "only_magic")]
#[test_case(Path::new("code_and_magic.json"), true; "code_and_magic")]
#[test_case(Path::new("only_code.json"), true; "only_code")]
#[test_case(Path::new("cell_magic.json"), false; "cell_magic")]
#[test_case(Path::new("valid_cell_magic.json"), true; "valid_cell_magic")]
#[test_case(Path::new("automagic.json"), false; "automagic")]
#[test_case(Path::new("automagics.json"), false; "automagics")]
#[test_case(Path::new("automagic_before_code.json"), false; "automagic_before_code")]
#[test_case(Path::new("automagic_after_code.json"), true; "automagic_after_code")]
fn test_is_valid_code_cell(path: &Path, expected: bool) -> Result<()> {
#[test_case("markdown", false)]
#[test_case("only_magic", true)]
#[test_case("code_and_magic", true)]
#[test_case("only_code", true)]
#[test_case("cell_magic", false)]
#[test_case("valid_cell_magic", true)]
#[test_case("automagic", false)]
#[test_case("automagics", false)]
#[test_case("automagic_before_code", false)]
#[test_case("automagic_after_code", true)]
#[test_case("unicode_magic_gh9145", true)]
fn test_is_valid_code_cell(cell: &str, expected: bool) -> Result<()> {
/// Read a Jupyter cell from the `resources/test/fixtures/jupyter/cell` directory.
fn read_jupyter_cell(path: impl AsRef<Path>) -> Result<Cell> {
let path = notebook_path("cell").join(path);
let source_code = std::fs::read_to_string(path)?;
Ok(serde_json::from_str(&source_code)?)
}

assert_eq!(read_jupyter_cell(path)?.is_valid_code_cell(), expected);
assert_eq!(
read_jupyter_cell(format!("{cell}.json"))?.is_valid_code_cell(),
expected
);
Ok(())
}

Expand Down

0 comments on commit cd3c2f7

Please sign in to comment.