Repository: ayasyrev/nbmetaclean (language: python)

Issue:
Create .pre-commit-hooks.yaml



Discussion / hints:

Fix patch:
diff --git a/.flake8 b/.flake8
index 99df13e..6ed7291 100644
--- a/.flake8
+++ b/.flake8
@@ -2,6 +2,7 @@
 select = C,E,F,W
 max-complexity = 10
 max-line-length = 120
+extend-ignore = W503
 disable-noqa = True
 application-import-names = nbmetaclean, tests
 import-order-style = google
diff --git a/.pre-commit-hooks.yaml b/.pre-commit-hooks.yaml
new file mode 100644
index 0000000..6fa2a48
--- /dev/null
+++ b/.pre-commit-hooks.yaml
@@ -0,0 +1,7 @@
+- id: nbclean
+  name: nbclean
+  description: Clean Jipyter Notebooks metadata and optionally output.
+  entry: nbclean
+  files: \.ipynb$
+  language: python
+  language_version: python3
diff --git a/noxfile.py b/noxfile.py
new file mode 100644
index 0000000..770d95b
--- /dev/null
+++ b/noxfile.py
@@ -0,0 +1,9 @@
+import nox
+
+
+@nox.session(python=["3.8", "3.9", "3.10", "3.11", "3.12"])
+def tests(session: nox.Session) -> None:
+    args = session.posargs or ["--cov"]
+    session.install("-r", "requirements_test.txt")
+    session.install(".")
+    session.run("pytest", *args)
diff --git a/noxfile_conda.py b/noxfile_conda.py
new file mode 100644
index 0000000..7ba7b37
--- /dev/null
+++ b/noxfile_conda.py
@@ -0,0 +1,9 @@
+import nox
+
+
+@nox.session(python=["3.8", "3.9", "3.10", "3.11", "3.12"], venv_backend="mamba")
+def conda_tests(session: nox.Session) -> None:
+    args = session.posargs or ["--cov"]
+    session.conda_install("--file", "requirements_test.txt")
+    session.install(".")
+    session.run("pytest", *args)
diff --git a/noxfile_conda_lint.py b/noxfile_conda_lint.py
new file mode 100644
index 0000000..f4052aa
--- /dev/null
+++ b/noxfile_conda_lint.py
@@ -0,0 +1,10 @@
+import nox
+
+locations = "src/nbmetaclean", "tests", "noxfile.py"
+
+
+@nox.session(python=["3.8", "3.9", "3.10", "3.11", "3.12"], venv_backend="mamba")
+def conda_lint(session: nox.Session) -> None:
+    args = session.posargs or locations
+    session.conda_install("flake8")
+    session.run("flake8", *args)
diff --git a/noxfile_lint.py b/noxfile_lint.py
new file mode 100644
index 0000000..db17853
--- /dev/null
+++ b/noxfile_lint.py
@@ -0,0 +1,11 @@
+import nox
+
+
+locations = "src/nbmetaclean", "tests", "noxfile.py"
+
+
+@nox.session(python=["3.8", "3.9", "3.10", "3.11", "3.12"])
+def lint(session: nox.Session) -> None:
+    args = session.posargs or locations
+    session.install("flake8")
+    session.run("flake8", *args)
diff --git a/requirements.txt b/requirements.txt
index e69de29..2522e8e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -0,0 +1,2 @@
+nbformat
+rich
\ No newline at end of file
diff --git a/setup.cfg b/setup.cfg
index c4d965d..4d6c8b7 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -13,6 +13,7 @@ classifiers =
     Programming Language :: Python :: 3.9
     Programming Language :: Python :: 3.10
     Programming Language :: Python :: 3.11
+    Programming Language :: Python :: 3.12
     License :: OSI Approved :: Apache Software License
     Operating System :: OS Independent
 
@@ -24,3 +25,9 @@ python_requires = >=3.8
 
 [options.packages.find]
 where = src
+
+[options.entry_points]
+console_scripts =
+    nbclean=nbmetaclean.app:app
+pipx.run =
+    nbclean=nbmetaclean.app:app
\ No newline at end of file
diff --git a/src/nbmetaclean/app.py b/src/nbmetaclean/app.py
new file mode 100644
index 0000000..16649c8
--- /dev/null
+++ b/src/nbmetaclean/app.py
@@ -0,0 +1,46 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+import nbformat
+from argparsecfg.app import App
+from argparsecfg.core import field_argument
+
+from .core import get_nb_names
+from .clean import clean_nb_file
+
+
+@dataclass
+class AppCfg:
+    path: str = field_argument("path", default=".", nargs="*")
+    as_version: int = field_argument(
+        default=nbformat.NO_CONVERT,
+        help="Save as version, default - no convert"
+    )
+    silent: bool = field_argument("-s", default=False, action="store_true")
+
+
+app = App(
+    prog="nbclean",
+    description="Clean metadata and execution_count from Jupyter notebooks.",
+)
+
+
+@app.main
+def clean(
+    cfg: AppCfg,
+) -> None:
+    """Clean metadata and execution_count from Jupyter notebook."""
+    path_list = cfg.path if isinstance(cfg.path, list) else [cfg.path]
+    nb_files: list[Path] = []
+    for path in path_list:
+        try:
+            nb_files.extend(get_nb_names(path))
+        except FileNotFoundError:
+            print(f"{path} not exists!")
+    print(f"find notebooks: {len(nb_files)} ")
+    cleaned = clean_nb_file(nb_files, as_version=cfg.as_version, silent=cfg.silent)
+    print(f"cleaned nbs: {len(cleaned)}")
+
+
+if __name__ == "__main__":
+    app()
diff --git a/src/nbmetaclean/clean.py b/src/nbmetaclean/clean.py
new file mode 100644
index 0000000..90d4236
--- /dev/null
+++ b/src/nbmetaclean/clean.py
@@ -0,0 +1,165 @@
+from __future__ import annotations
+import copy
+
+from pathlib import Path
+from typing import Optional, Union
+
+import nbformat
+
+from nbformat.notebooknode import NotebookNode
+from rich.progress import track
+
+from .core import read_nb, write_nb, PathOrStr
+
+
+NB_METADATA_PRESERVE_MASKS = [
+    ("language_info", "name"),
+]
+
+
+def get_meta_by_mask(
+    nb_meta: Union[NotebookNode, str],
+    mask: Optional[tuple[str, ...]] = None,
+) -> Union[NotebookNode, str]:
+    """Get metadata by mask."""
+    if isinstance(nb_meta, str) or mask == ():
+        return nb_meta
+    if mask is None:
+        return NotebookNode()
+    result = {}
+    value = nb_meta.get(mask[0])
+    if value is not None:
+        new_mask = tuple(mask[1:])
+        result[mask[0]] = get_meta_by_mask(value, new_mask) or value
+    return NotebookNode(result)
+
+
+def new_metadata(
+    nb_meta: NotebookNode,
+    masks: Optional[list[tuple[str, ...]]] = None,
+) -> NotebookNode:
+    """Clean notebooknode metadata."""
+    if masks is None:
+        return NotebookNode()
+    filtered_meta = {}
+    for mask in masks:
+        filtered_meta.update(get_meta_by_mask(nb_meta, mask))
+    return NotebookNode(filtered_meta)
+
+
+def clean_cell_metadata(
+    cell: NotebookNode,
+    clear_execution_count: bool = True,
+    clear_outputs: bool = False,
+    preserve_cell_metadata_mask: Optional[list[tuple[str, ...]]] = None,
+) -> bool:
+    """Clean cell metadata."""
+    changed = False
+    if cell.get("metadata", None):
+        old_metadata = copy.deepcopy(cell.metadata)
+        cell.metadata = new_metadata(cell.metadata, preserve_cell_metadata_mask)
+        if cell.metadata != old_metadata:
+            changed = True
+    if clear_outputs and hasattr(cell, "outputs") and cell.get("outputs", None):
+        cell.outputs = []
+        changed = True
+    if (
+        clear_execution_count
+        and hasattr(cell, "execution_count")
+        and cell.execution_count
+    ):
+        cell.execution_count = None
+        changed = True
+    if hasattr(cell, "outputs") and cell.outputs:
+        for output in cell.outputs:
+            if clear_execution_count and output.get("execution_count", None):
+                output.execution_count = None
+                changed = True
+            if output.get("metadata", None):
+                old_metadata = copy.deepcopy(output.metadata)
+                output.metadata = new_metadata(
+                    output.metadata, preserve_cell_metadata_mask
+                )
+                if output.metadata != old_metadata:
+                    changed = True
+    return changed
+
+
+def clean_nb(
+    nb: NotebookNode,
+    clear_nb_metadata: bool = True,
+    clear_cell_metadata: bool = True,
+    clear_execution_count: bool = True,
+    clear_outputs: bool = False,
+    preserve_nb_metadata_masks: Optional[list[tuple[str, str]],] = None,
+    preserve_cell_metadata_mask: Optional[str] = None,
+) -> tuple[NotebookNode, bool]:
+    """Clean notebook - metadata, execution_count, outputs.
+
+    Args:
+        nb (Notebook): Notebook to clean.
+        clear_execution_count (bool, optional): Clear execution_count. Defaults to True.
+        clear_outputs (bool, optional): Clear outputs. Defaults to False.
+
+    Returns:
+        bool: True if changed.
+    """
+    changed = False
+    if clear_nb_metadata and nb.metadata:
+        old_metadata = copy.deepcopy(nb.metadata)
+        masks = preserve_nb_metadata_masks or NB_METADATA_PRESERVE_MASKS
+        nb.metadata = new_metadata(nb.metadata, masks=masks)
+        if nb.metadata != old_metadata:
+            changed = True
+    if clear_cell_metadata:
+        for cell in nb.cells:
+            result = clean_cell_metadata(
+                cell,
+                clear_execution_count=clear_execution_count,
+                clear_outputs=clear_outputs,
+                preserve_cell_metadata_mask=preserve_cell_metadata_mask,
+            )
+            if result:
+                changed = True
+
+    return nb, changed
+
+
+def clean_nb_file(
+    path: Union[PathOrStr, list[PathOrStr]],
+    clear_nb_metadata: bool = True,
+    clear_cell_metadata: bool = True,
+    clear_execution_count: bool = True,
+    clear_outputs: bool = False,
+    as_version: nbformat.Sentinel = nbformat.NO_CONVERT,
+    silent: bool = False,
+) -> list[Path]:
+    """Clean metadata and execution count from notebook.
+
+    Args:
+        path (Union[str, PosixPath]): Notebook filename or list of names.
+        as_version (int, optional): Nbformat version. Defaults to 4.
+        clear_execution_count (bool, optional): Clean execution count. Defaults to True.
+        silent (bool, optional): Silent mode. Defaults to False.
+
+    Returns:
+        List[Path]: List of cleaned notebooks
+    """
+    if not isinstance(path, list):
+        path = [path]
+    cleaned: list[PathOrStr] = []
+    for filename in track(path, transient=True):
+        nb = read_nb(filename)
+        nb, result = clean_nb(
+            nb,
+            clear_execution_count=clear_execution_count,
+            clear_outputs=clear_outputs,
+            clear_nb_metadata=clear_nb_metadata,
+            clear_cell_metadata=clear_cell_metadata,
+        )
+        if result:
+            cleaned.append(filename)
+            write_nb(nb, filename, as_version)
+            if not silent:
+                print(f"done: {filename}")
+    return cleaned
diff --git a/src/nbmetaclean/core.py b/src/nbmetaclean/core.py
new file mode 100644
index 0000000..64c73cc
--- /dev/null
+++ b/src/nbmetaclean/core.py
@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from pathlib import Path, PosixPath
+from typing import Optional, TypeVar
+
+import nbformat
+from nbformat.notebooknode import NotebookNode
+
+PathOrStr = TypeVar("PathOrStr", str, Path, PosixPath)
+
+
+def read_nb(path: PathOrStr) -> NotebookNode:
+    """Read notebook from filename.
+
+    Args:
+        path (Union[str, PosixPath): Notebook filename.
+
+    Returns:
+        Notebook: Jupyter Notebook.
+    """
+    with Path(path).open("r", encoding="utf-8") as fh:
+        nb: NotebookNode = nbformat.read(fh, as_version=nbformat.NO_CONVERT)
+    return nb
+
+
+def write_nb(
+    nb: NotebookNode,
+    path: PathOrStr,
+    as_version: nbformat.Sentinel = nbformat.NO_CONVERT,
+) -> Path:
+    """Write notebook to file
+
+    Args:
+        nb (Notebook): Notebook to write
+        path (Union[str, PosixPath]): filename to write
+        as_version (_type_, optional): Nbformat version. Defaults to nbformat.NO_CONVERT.
+    Returns:
+        Path: Filename of writed Nb.
+    """
+    filename = Path(path)
+    if filename.suffix != ".ipynb":
+        filename = filename.with_suffix(".ipynb")
+    with filename.open("w", encoding="utf-8") as fh:
+        nbformat.write(nb, fh, version=as_version)  # type: ignore
+    return filename
+
+
+def get_nb_names(
+        path: Optional[PathOrStr] = None,
+        recursive: bool = True,
+        filter_hidden: bool = True,
+) -> list[Path]:
+    """Return list of notebooks from `path`. If no `path` return notebooks from current folder.
+
+    Args:
+        path (Union[Path, str, None]): Path for nb or folder with notebooks.
+        recursive bool: Recursive search.
+        filter_hidden bool: Filter hidden paths.
+
+    Raises:
+        sys.exit: If filename or dir not exists or not nb file.
+
+    Returns:
+        List[Path]: List of notebooks names.
+    """
+    nb_path = Path(path or ".")
+
+    if not nb_path.exists():
+        raise FileNotFoundError(f"{nb_path} not exists!")
+
+    if nb_path.is_dir():
+        result = []
+        for item in nb_path.iterdir():
+            if item.is_file() and item.suffix == ".ipynb":
+                if filter_hidden and item.name.startswith("."):
+                    continue
+                result.append(item)
+            if item.is_dir():
+                if recursive:
+                    if filter_hidden and item.name.startswith("."):
+                        continue
+                    result.extend(get_nb_names(item, recursive, filter_hidden))
+
+        return result
+
+    return [nb_path]
diff --git a/src/nbmetaclean/version.py b/src/nbmetaclean/version.py
index f102a9c..1e09a73 100644
--- a/src/nbmetaclean/version.py
+++ b/src/nbmetaclean/version.py
@@ -1,1 +1,1 @@
-__version__ = "0.0.1"
+__version__ = "0.0.1"  # pragma: no cover


Test patch:
diff --git a/requirements_test.txt b/requirements_test.txt
index 036d8c5..16fc33b 100644
--- a/requirements_test.txt
+++ b/requirements_test.txt
@@ -1,2 +1,3 @@
 pytest
-flake8
\ No newline at end of file
+flake8
+pytest-cov
diff --git a/requirements_test_extra.txt b/requirements_test_extra.txt
index 50ee489..7524e2e 100644
--- a/requirements_test_extra.txt
+++ b/requirements_test_extra.txt
@@ -1,3 +1,3 @@
-pytest-cov
 coverage[toml]
-nox
\ No newline at end of file
+nox
+black
diff --git a/tests/test_clean.py b/tests/test_clean.py
new file mode 100644
index 0000000..e194548
--- /dev/null
+++ b/tests/test_clean.py
@@ -0,0 +1,197 @@
+import copy
+from pathlib import Path
+from nbformat.notebooknode import NotebookNode
+
+from pytest import CaptureFixture
+
+from nbmetaclean.clean import (
+    NB_METADATA_PRESERVE_MASKS,
+    clean_cell_metadata,
+    clean_nb,
+    clean_nb_file,
+    get_meta_by_mask,
+    new_metadata,
+)
+from nbmetaclean.core import read_nb, write_nb
+
+
+def test_get_meta_by_mask():
+    """test get_meta_by_mask"""
+    nb = read_nb(Path("tests/test_nbs/test_nb_2.ipynb"))
+    nb_meta = nb.metadata
+
+    # string as nb_meta
+    new_meta = get_meta_by_mask("some string")
+    assert new_meta == "some string"
+
+    # no mask
+    new_meta = get_meta_by_mask(nb_meta)
+    assert new_meta == {}
+
+    # mask
+    nb_meta["some key"] = "some value"
+    new_meta = get_meta_by_mask(nb_meta, ("some key",))
+    assert new_meta == {"some key": "some value"}
+    new_meta = get_meta_by_mask(nb_meta, NB_METADATA_PRESERVE_MASKS[0])
+    assert new_meta == {"language_info": {"name": "python"}}
+
+    # mask for empty result
+    new_meta = get_meta_by_mask(nb_meta, ("some other key",))
+    assert new_meta == {}
+
+
+def test_new_metadata():
+    """test new_metadata"""
+    nb_meta = read_nb("tests/test_nbs/test_nb_2.ipynb").metadata
+    new_meta = new_metadata(nb_meta)
+    assert isinstance(new_meta, NotebookNode)
+    assert not new_meta
+    new_meta = new_metadata(nb_meta, [("language_info", "name")])
+    assert new_meta == {"language_info": {"name": "python"}}
+
+
+def test_clean_cell_metadata():
+    """test clean_cell_metadata"""
+    test_nb = read_nb("tests/test_nbs/test_nb_2.ipynb")
+
+    # clear outputs
+    cell = copy.deepcopy(test_nb.cells[1])
+    assert cell.outputs
+    assert not cell.metadata
+    assert cell.execution_count == 1
+    cell.metadata = {"some key": "some value"}
+    changed = clean_cell_metadata(cell, clear_outputs=True)
+    assert changed
+    assert not cell.outputs
+    assert not cell.metadata
+    assert not cell.execution_count
+
+    # dont clear outputs, execution_count, mask
+    cell = copy.deepcopy(test_nb.cells[1])
+    cell.metadata = {"some key": "some value"}
+    cell.outputs[0].metadata = {
+        "some key": "some value",
+        "some other key": "some value",
+    }
+    changed = clean_cell_metadata(
+        cell,
+        clear_execution_count=False,
+        preserve_cell_metadata_mask=[("some key",)],
+    )
+    assert changed
+    assert cell.outputs[0].metadata == {"some key": "some value"}
+    assert cell.metadata == {"some key": "some value"}
+    assert cell.execution_count == 1
+
+    # clear outputs, same mask -> no changes meta, clear ex
+    changed = clean_cell_metadata(cell, preserve_cell_metadata_mask=[("some key",)])
+    assert changed
+    assert cell.execution_count is None
+    assert cell.metadata == {"some key": "some value"}
+
+    # clear execution_count, metadata
+    changed = clean_cell_metadata(cell)
+    assert changed
+    assert not cell.outputs[0].metadata
+    assert not cell.execution_count
+    assert not cell.metadata
+    assert not cell.outputs[0].metadata
+
+
+def test_clean_cell_metadata_markdown():
+    """test clean_cell_metadata with markdown cell"""
+    test_nb = read_nb("tests/test_nbs/test_nb_2.ipynb")
+    cell = copy.deepcopy(test_nb.cells[0])
+    cell.metadata = {"some key": "some value"}
+    changed = clean_cell_metadata(cell)
+    assert changed
+    assert not cell.metadata
+
+
+def test_clean_nb():
+    """test clean nb"""
+    path = Path("tests/test_nbs")
+    nb_path = path / "test_nb_2.ipynb"
+    nb_clean = path / "test_nb_2_clean.ipynb"
+    nb = read_nb(nb_path)
+    assert nb.cells[1].execution_count == 1
+    assert nb.cells[1].outputs[0].execution_count == 1
+    assert nb.metadata
+    nb, result = clean_nb(nb)
+    assert result is True
+    assert nb.cells[1].execution_count is None
+    assert nb.cells[1].outputs[0].execution_count is None
+    nb_clean = read_nb(nb_clean)
+    assert nb == nb_clean
+
+    # # try clean cleaned
+    nb, result = clean_nb(nb_clean)
+    assert not result
+
+    # # clean metadata, leave execution_count
+    nb = read_nb(nb_path)
+    nb, result = clean_nb(nb, clear_execution_count=False)
+    assert result
+    assert nb.cells[1].execution_count == 1
+    assert nb.cells[1].outputs[0].execution_count == 1
+    assert nb.metadata == nb_clean.metadata
+
+    # clean nb metadata, leave cells metadata
+    nb = read_nb(nb_path)
+    nb.cells[1].metadata = {"some key": "some value"}
+    nb, result = clean_nb(nb, clear_cell_metadata=False)
+    assert result
+    assert nb.metadata == nb_clean.metadata
+    assert nb.cells[1].metadata == {"some key": "some value"}
+    assert nb.cells[1].execution_count == 1
+
+    # clean cells metadata, leave nb metadata
+    nb = read_nb(nb_path)
+    nb_meta = copy.deepcopy(nb.metadata)
+    nb, result = clean_nb(nb, clear_nb_metadata=False)
+    assert result
+    assert nb.metadata == nb_meta
+    assert nb.cells[1].execution_count is None
+
+
+def test_clean_nb_file(tmp_path: Path, capsys: CaptureFixture[str]):
+    """test clean nb file"""
+    path = Path("tests/test_nbs")
+    nb_name = "test_nb_2.ipynb"
+    nb_clean = read_nb(path / "test_nb_2_clean.ipynb")
+
+    # prepare temp test notebook
+    test_nb_path = write_nb(read_nb(path / nb_name), tmp_path / nb_name)
+
+    # clean meta, leave execution_count
+    result = clean_nb_file(test_nb_path, clear_execution_count=False)
+    assert len(result) == 1
+    nb = read_nb(result[0])
+    assert nb.metadata == nb_clean.metadata
+    assert nb.cells[1].execution_count == 1
+    assert nb.cells[1].outputs[0].execution_count == 1
+
+    # clean meta, execution_count
+    # path as list
+    result = clean_nb_file([test_nb_path])
+    captured = capsys.readouterr()
+    out = captured.out
+    assert out.startswith("done:")
+    assert "test_clean_nb_file0/test_nb_2.ipynb" in out
+    assert len(result) == 1
+    nb = read_nb(result[0])
+    assert nb == nb_clean
+
+    # try clean cleaned
+    result = clean_nb_file(test_nb_path)
+    assert len(result) == 0
+    captured = capsys.readouterr()
+    out = captured.out
+    assert not out.strip()
+
+    # silent
+    test_nb_path = write_nb(read_nb(path / nb_name), tmp_path / nb_name)
+    result = clean_nb_file(test_nb_path, silent=True)
+    # assert len(result) == 1
+    captured = capsys.readouterr()
+    assert not captured.out.strip()
diff --git a/tests/test_nbs/test_nb_1.ipynb b/tests/test_nbs/test_nb_1.ipynb
new file mode 100644
index 0000000..5c6deca
--- /dev/null
+++ b/tests/test_nbs/test_nb_1.ipynb
@@ -0,0 +1,23 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/tests/test_nbs/test_nb_2.ipynb b/tests/test_nbs/test_nb_2.ipynb
new file mode 100644
index 0000000..b75ce89
--- /dev/null
+++ b/tests/test_nbs/test_nb_2.ipynb
@@ -0,0 +1,52 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "markdown cell source"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "2"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "1 + 1"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "nbmetaclean",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/tests/test_nbs/test_nb_2_clean.ipynb b/tests/test_nbs/test_nb_2_clean.ipynb
new file mode 100644
index 0000000..f316a23
--- /dev/null
+++ b/tests/test_nbs/test_nb_2_clean.ipynb
@@ -0,0 +1,38 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "markdown cell source"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "2"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "1 + 1"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/tests/test_read_write.py b/tests/test_read_write.py
new file mode 100644
index 0000000..778348e
--- /dev/null
+++ b/tests/test_read_write.py
@@ -0,0 +1,62 @@
+from pathlib import Path
+
+from nbmetaclean.core import get_nb_names, read_nb, write_nb
+
+
+def test_read_nb():
+    """test read notebook"""
+    file = Path("tests/test_nbs/test_nb_1.ipynb")
+    nb = read_nb(file)
+    assert isinstance(nb, dict)
+    assert nb["metadata"] == {"language_info": {"name": "python"}}
+    assert nb["nbformat"] == 4
+    assert nb["nbformat_minor"] == 2
+    cells = nb["cells"]
+    assert isinstance(cells, list)
+    assert len(cells) == 2
+    # markdown
+    assert cells[0]["cell_type"] == "markdown"
+    assert cells[0]["source"] == ""
+    assert cells[0]["metadata"] == {}
+    # code
+    assert cells[1]["cell_type"] == "code"
+    assert cells[1]["source"] == ""
+    assert cells[1]["execution_count"] is None
+    assert cells[1]["metadata"] == {}
+    assert cells[1]["outputs"] == []
+
+
+def test_write_nb(tmp_path: Path):
+    """test write notebook"""
+    file = Path("tests/test_nbs/test_nb_1.ipynb")
+    nb = read_nb(file)
+    write_nb(nb, tmp_path / file.name)
+    with open(tmp_path / file.name, "r", encoding="utf-8") as fh:
+        res_text = fh.read()
+    with open(file, "r", encoding="utf-8") as fh:
+        org_text = fh.read()
+    assert res_text == org_text
+
+    # write with name w/o suffix
+    result = write_nb(nb, tmp_path / "test_nb_1")
+    assert result == tmp_path / "test_nb_1.ipynb"
+
+
+def test_get_nb_names():
+    """test get_nb_names"""
+    path = Path("tests/test_nbs")
+    file = path / "test_nb_1.ipynb"
+    names = get_nb_names(file)
+    assert len(names) == 1
+    names.sort(key=lambda x: x.name)
+    assert names[0] == file
+    names = get_nb_names(path)
+    assert len(names) == 3
+    names.sort(key=lambda x: x.name)
+    assert names[0] == file
+    try:
+        get_nb_names("wrong_name")
+        assert False
+    except FileNotFoundError as ex:
+        assert True
+        assert str(ex) == "wrong_name not exists!"