Add Jupyter notebook support (#2)

Author: cyyc1

.flake8

@@ -2,3 +2,4 @@
 max-complexity = 10
 max-line-length = 88
 ignore = E203, W503
+exclude = tests/resources/*

.pre-commit-config.yaml

@@ -1,3 +1,4 @@
+exclude: '^tests/resources/.*$'  # ignore files in that folder because they are designed that way
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.3.0

.pre-commit-hooks.yaml

@@ -1,6 +1,12 @@
 -   id: format-def-indent
-    name: Format indentation in function definitions
+    name: Format indentation in function definitions (in .py files)
     description: Make sure arguments in Python function definitions are indented by 8 spaces.
     entry: format-def-indent
     language: python
     types: [python]
+-   id: format-def-indent-in-jupyter
+    name: Format indentation in function definitions (in .ipynb files)
+    description: In Jupyter notebooks, make sure arguments in Python function definitions are indented by 8 spaces.
+    entry: format-def-indent-in-jupyter
+    language: python
+    files: '^.*\.ipynb$'

README.md

@@ -20,20 +20,37 @@ pip install format-def-indent
 
 ### As a command line tool
 
+To auto-format Python files (.py):
+
 ```bash
 format-def-indent <PATH_THAT_CONTAINS_PYTHON_FILES>
 ```
 Use `--help` to see documentations of command line arguments.
 
+To auto-format Jupyter notebooks (.ipynb):
+
+```bash
+format-def-indent-in-jupyter <PATH_THAT_CONTAINS_PYTHON_FILES>
+```
+
 ### As a pre-commit hook
 
-Put the following into your `.pre-commit-config.yaml` file. Remember to replace `<VERSION>` with your version of this tool (such as `v0.1.4`):
+To auto-format Python files (.py), put the following into your `.pre-commit-config.yaml` file. Remember to replace `<VERSION>` with your version of this tool (such as `v0.1.4`):
 ```yaml
 -   repo: https://github.com/cyyc1/py-def-indent-formatter
     rev: <VERSION>
     hooks:
     -   id: format-def-indent
 ```
+
+To auto-format Jupyter notebooks (.ipynb), put the following into your `.pre-commit-config.yaml` file:
+```yaml
+-   repo: https://github.com/cyyc1/py-def-indent-formatter
+    rev: <VERSION>
+    hooks:
+    -   id: format-def-indent-in-jupyter
+```
+
 See [pre-commit](https://github.com/pre-commit/pre-commit) for more instructions.
 
 ## What does this formatter do

format_def_indent/__main__.py

@@ -0,0 +1,26 @@
+import argparse
+
+from pathlib import Path
+
+
+class BaseFixer:
+    def __init__(self, path: str, cli_args: argparse.Namespace) -> None:
+        self.path = path
+        self.cli_args = cli_args
+
+    def fix_one_directory_or_one_file(self) -> int:
+        path_obj = Path(self.path)
+
+        if path_obj.is_file():
+            return self.fix_one_file(path_obj.as_posix())
+
+        filenames = sorted(path_obj.rglob('*.py'))
+        all_status = set()
+        for filename in filenames:
+            status = self.fix_one_file(filename)
+            all_status.add(status)
+
+        return 0 if all_status == {0} else 1
+
+    def fix_one_file(self, *varargs, **kwargs):
+        raise NotImplementedError('Please implement this method')

format_def_indent/_base_fixer.py

@@ -0,0 +1,69 @@
+import argparse
+import json
+import sys
+from typing import Sequence, Optional
+
+from jupyter_notebook_parser import JupyterNotebookParser, JupyterNotebookRewriter
+
+import format_def_indent._helper as helper
+from format_def_indent._base_fixer import BaseFixer
+
+
+class JupyterNotebookFixer(BaseFixer):
+    def __init__(self, path: str, cli_args: argparse.Namespace) -> None:
+        super().__init__(path=path, cli_args=cli_args)
+
+    def fix_one_file(self, filename: str) -> int:
+        try:
+            parsed = JupyterNotebookParser(filename)
+            rewriter = JupyterNotebookRewriter(parsed_notebook=parsed)
+            code_cells = parsed.get_code_cells()
+            code_cell_indices = parsed.get_code_cell_indices()
+            code_cell_sources = parsed.get_code_cell_sources()
+        except Exception as exc:
+            print(f'Error reading {filename}: {str(exc)}', file=sys.stderr)
+            return 1
+        else:
+            ret_val = 0
+            assert len(code_cells) == len(code_cell_indices)
+            assert len(code_cells) == len(code_cell_sources)
+
+            for i in range(len(code_cells)):
+                index: int = code_cell_indices[i]
+                source: str = code_cell_sources[i]
+                fixed: str = helper.fix_src(source_code=source)
+
+                if fixed != source:
+                    ret_val = 1
+                    rewriter.replace_source_in_code_cell(
+                        index=index,
+                        new_source=fixed,
+                    )
+
+            if ret_val == 1:
+                print(f'Rewriting {filename}', file=sys.stderr)
+                with open(filename, 'w') as fp:
+                    json.dump(parsed.notebook_content, fp, indent=1)
+                    # Jupyter notebooks (.ipynb) always ends with a new line
+                    # but json.dump does not.
+                    fp.write('\n')
+
+            return 0 if self.cli_args.exit_zero_even_if_changed else ret_val
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('paths', nargs='*')
+    parser.add_argument('--exit-zero-even-if-changed', action='store_true')
+    args = parser.parse_args(argv)
+
+    ret = 0
+    for path in args.paths:
+        fixer = JupyterNotebookFixer(path=path, cli_args=args)
+        ret |= fixer.fix_one_directory_or_one_file()
+
+    return ret
+
+
+if __name__ == '__main__':
+    raise SystemExit(main())

format_def_indent/_main.py

@@ -0,0 +1,60 @@
+# This file is inspired by https://github.com/asottile/add-trailing-comma/blob/6be6dfc05176bddfc5a9c4bf0fd4941850f0fb41/add_trailing_comma/_main.py  # noqa: E501
+
+import argparse
+import sys
+
+from typing import Sequence, Optional
+
+import format_def_indent._helper as helper
+from format_def_indent._base_fixer import BaseFixer
+
+
+class PythonFileFixer(BaseFixer):
+    def __init__(self, path: str, cli_args: argparse.Namespace) -> None:
+        super().__init__(path=path, cli_args=cli_args)
+
+    def fix_one_file(self, filename: str) -> int:
+        if filename == '-':
+            source_bytes = sys.stdin.buffer.read()
+        else:
+            with open(filename, 'rb') as fb:
+                source_bytes = fb.read()
+
+        try:
+            source_text_orig = source_text = source_bytes.decode()
+        except UnicodeDecodeError:
+            msg = f'{filename} is non-utf-8 (not supported)'
+            print(msg, file=sys.stderr)
+            return 1
+
+        source_text = helper.fix_src(source_text)
+
+        if filename == '-':
+            print(source_text, end='')
+        elif source_text != source_text_orig:
+            print(f'Rewriting {filename}', file=sys.stderr)
+            with open(filename, 'wb') as f:
+                f.write(source_text.encode())
+
+        if self.cli_args.exit_zero_even_if_changed:
+            return 0
+        else:
+            return source_text != source_text_orig
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('paths', nargs='*')
+    parser.add_argument('--exit-zero-even-if-changed', action='store_true')
+    args = parser.parse_args(argv)
+
+    ret = 0
+    for path in args.paths:
+        fixer = PythonFileFixer(path=path, cli_args=args)
+        ret |= fixer.fix_one_directory_or_one_file()
+
+    return ret
+
+
+if __name__ == '__main__':
+    raise SystemExit(main())

format_def_indent/_main_jupyter.py

@@ -1,6 +1,6 @@
 [metadata]
 name = format_def_indent
-version = 0.1.4
+version = 0.1.5
 description = Make sure arguments in Python function definitions are indented by 8 spaces
 long_description = file: README.md
 long_description_content_type = text/markdown
@@ -15,6 +15,7 @@ classifiers =
 [options]
 packages = find:
 install_requires =
+    jupyter-notebook-parser>=0.1.3
     tokenize-rt>=3.0.1
 python_requires = >=3.8
 
@@ -25,7 +26,8 @@ exclude =
 
 [options.entry_points]
 console_scripts =
-    format-def-indent = format_def_indent._main:main
+    format-def-indent = format_def_indent._main_py:main
+    format-def-indent-in-jupyter = format_def_indent._main_jupyter:main
 
 [bdist_wheel]
 universal = True