RustPython · youknowone · Jan 17, 2026 · Jan 16, 2026 · Jan 16, 2026 · Jan 16, 2026
diff --git a/.claude/commands/upgrade-pylib.md b/.claude/commands/upgrade-pylib.md
@@ -15,9 +15,62 @@ Upgrade a Python standard library module from CPython to RustPython.
    - If `cpython/Lib/$ARGUMENTS.py` exists, copy it to `Lib/$ARGUMENTS.py`
    - If `cpython/Lib/$ARGUMENTS/` directory exists, copy it to `Lib/$ARGUMENTS/`
 
-3. **Upgrade tests**
-   - Run: `python lib_updater.py --quick-upgrade cpython/Lib/test/test_$ARGUMENTS`
-   - This will update the test files with appropriate RustPython markers
+3. **Upgrade tests (quick upgrade with lib_updater)**
+   - If `cpython/Lib/test/test_$ARGUMENTS.py` is a single file:
+     - Run: `python3 scripts/lib_updater.py --quick-upgrade cpython/Lib/test/test_$ARGUMENTS.py`
+   - If `cpython/Lib/test/test_$ARGUMENTS/` is a directory:
+     - Run the script for each `.py` file in the directory:
+       ```bash
+       for f in cpython/Lib/test/test_$ARGUMENTS/*.py; do
+           python3 scripts/lib_updater.py --quick-upgrade "$f"
+       done
+       ```
+   - This will update the test files with basic RustPython markers (`@unittest.expectedFailure`, `@unittest.skip`, etc.)
+   - **Handle lib_updater warnings**: If you see warnings like `WARNING: TestCFoo does not exist in remote file`, it means the class structure changed between versions and markers couldn't be transferred automatically. These need to be manually restored in step 4 or added in step 5.
+
+4. **Review git diff and restore RUSTPYTHON-specific changes**
+   - Run `git diff Lib/test/test_$ARGUMENTS` to review all changes
+   - **Only restore changes that have explicit `RUSTPYTHON` comments**. Look for:
+     - `# XXX: RUSTPYTHON` or `# XXX RUSTPYTHON` - Comments marking RustPython-specific code modifications
+     - `# TODO: RUSTPYTHON` - Comments marking tests that need work
+     - Code changes with inline `# ... RUSTPYTHON` comments
+   - **Do NOT restore other diff changes** - these are likely upstream CPython changes, not RustPython-specific modifications
+   - When restoring, preserve the original context and formatting
+
+5. **Verify tests**
+   - Run: `cargo run --release -- -m test test_$ARGUMENTS -v`
+   - The `-v` flag shows detailed output to identify which tests fail and why
+   - For each new failure, add appropriate markers based on the failure type:
+     - **Test assertion failure** → `@unittest.expectedFailure` with `# TODO: RUSTPYTHON` comment
+     - **Panic/crash** → `@unittest.skip("TODO: RUSTPYTHON; <panic message>")`
+   - **Class-specific markers**: If a test fails only in the C implementation (TestCFoo) but passes in the Python implementation (TestPyFoo), or vice versa, add the marker to the specific subclass, not the base class:
+     ```python
+     # Base class - no marker here
+     class TestFoo:
+         def test_something(self):
+             ...
+
+     class TestPyFoo(TestFoo, PyTest): pass
+
+     class TestCFoo(TestFoo, CTest):
+         # TODO: RUSTPYTHON
+         @unittest.expectedFailure
+         def test_something(self):
+             return super().test_something()
+     ```
+   - **New tests from CPython**: The upgrade may bring in entirely new tests that didn't exist before. These won't have any RUSTPYTHON markers in the diff - they just need to be tested and marked if they fail.
+   - Example markers:
+     ```python
+     # TODO: RUSTPYTHON
+     @unittest.expectedFailure
+     def test_something(self):
+         ...
+
+     # TODO: RUSTPYTHON
+     @unittest.skip("TODO: RUSTPYTHON; panics with 'index out of bounds'")
+     def test_crashes(self):
+         ...
+     ```
 
 ## Example Usage
 ```
@@ -26,7 +79,30 @@ Upgrade a Python standard library module from CPython to RustPython.
 /upgrade-pylib asyncio
 ```
 
+## Example: Restoring RUSTPYTHON changes
+
+When git diff shows removed RUSTPYTHON-specific code like:
+```diff
+-# XXX RUSTPYTHON: we don't import _json as fresh since...
+-cjson = import_helper.import_fresh_module('json') #, fresh=['_json'])
++cjson = import_helper.import_fresh_module('json', fresh=['_json'])
+```
+
+You should restore the RustPython version:
+```python
+# XXX RUSTPYTHON: we don't import _json as fresh since...
+cjson = import_helper.import_fresh_module('json') #, fresh=['_json'])
+```
+
 ## Notes
 - The cpython/ directory should contain the CPython source that we're syncing from
-- lib_updater.py handles adding `# TODO: RUSTPYTHON` markers and `@unittest.expectedFailure` decorators
-- After upgrading, you may need to run tests to verify: `cargo run --release -- -m test test_$ARGUMENTS`
+- `scripts/lib_updater.py` handles basic patching:
+  - Transfers `@unittest.expectedFailure` and `@unittest.skip` decorators with `TODO: RUSTPYTHON` markers
+  - Adds `import unittest # XXX: RUSTPYTHON` if needed for the decorators
+  - **Limitation**: If a class was restructured (e.g., method overrides removed), lib_updater will warn and skip those markers
+- The script does NOT preserve all RustPython-specific changes - you must review `git diff` and restore them
+- Common RustPython markers to look for:
+  - `# XXX: RUSTPYTHON` or `# XXX RUSTPYTHON` - Inline comments for code modifications
+  - `# TODO: RUSTPYTHON` - Test skip/failure markers
+  - Any code with `RUSTPYTHON` in comments that was removed in the diff
+- **Important**: Not all changes in the git diff need to be restored. Only restore changes that have explicit `RUSTPYTHON` comments. Other changes are upstream CPython updates.
diff --git a/scripts/lib_updater.py b/scripts/lib_updater.py
@@ -270,11 +270,42 @@ def {test_name}(self):
             yield (lineno, textwrap.indent(patch_lines, DEFAULT_INDENT))
 
 
+def has_unittest_import(tree: ast.Module) -> bool:
+    """Check if 'import unittest' is already present in the file."""
+    for node in tree.body:
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                if alias.name == UT and alias.asname is None:
+                    return True
+    return False
+
+
+def find_import_insert_line(tree: ast.Module) -> int:
+    """Find the line number after the last import statement."""
+    last_import_line = None
+    for node in tree.body:
+        if isinstance(node, (ast.Import, ast.ImportFrom)):
+            last_import_line = node.end_lineno or node.lineno
+    assert last_import_line is not None
+    return last_import_line
+
+
 def apply_patches(contents: str, patches: Patches) -> str:
     tree = ast.parse(contents)
     lines = contents.splitlines()
 
     modifications = list(iter_patch_lines(tree, patches))
+
+    # If we have modifications and unittest is not imported, add it
+    if modifications and not has_unittest_import(tree):
+        import_line = find_import_insert_line(tree)
+        modifications.append(
+            (
+                import_line,
+                "\nimport unittest # XXX: RUSTPYTHON; importing to be able to skip tests",
+            )
+        )
+
     # Going in reverse to not distrupt the line offset
     for lineno, patch in sorted(modifications, reverse=True):
         lines.insert(lineno, patch)