Add main module

Author: andersy005

ci/environment-docs.yml

@@ -11,6 +11,11 @@ dependencies:
   - sphinx-autobuild
   - sphinx-copybutton
   - sphinx-inline-tabs
+  - rich
+  - typer
+  - jupyter_client
+  - nbformat
+  - nbconvert
   - pip:
       - sphinxext-opengraph
       - -e ..

ci/environment.yml

@@ -5,3 +5,8 @@ channels:
 dependencies:
   - pytest-cov
   - pre-commit
+  - rich
+  - typer
+  - jupyter_client
+  - nbformat
+  - nbconvert

ci/upstream-dev-environment.yml

@@ -5,3 +5,8 @@ channels:
 dependencies:
   - pytest-cov
   - pre-commit
+  - rich
+  - typer
+  - jupyter_client
+  - nbformat
+  - nbconvert

pyproject.toml

@@ -4,4 +4,4 @@ target-version = ['py38']
 skip-string-normalization = true
 
 [build-system]
-requires = ["setuptools>=45", "wheel", "setuptools_scm>=6.2"]
+requires = ["setuptools>=45", "wheel", "setuptools_scm>=6.2", "setuptools_scm[toml]>=6.2"]

requirements.txt

@@ -0,0 +1,4 @@
+rich>=11.2.0
+typer>=0.4.0
+nbformat
+jupyter-client

unison/console.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+import rich.console
+
+console = rich.console.Console()

unison/main.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import dataclasses
+import json
+import pathlib
+import subprocess
+
+import nbconvert.preprocessors
+import nbformat
+
+from .console import console
+
+
+@dataclasses.dataclass
+class Unison:
+    def get_conda_kernel_path(self, env_name) -> str | None:
+        command = ['conda', 'env', 'list', '--json']
+        try:
+            output = subprocess.check_output(command).decode('ascii')
+            envs = json.loads(output)['envs']
+
+            for env in envs:
+                env = pathlib.Path(env)
+                if env.stem == env_name:
+                    return env
+        except Exception:
+            pass
+
+        finally:
+            return None
+
+    def set_kernel_name(
+        self,
+        notebook_in: str | pathlib.Path,
+        kernel_name: str,
+        *,
+        notebook_out: str | pathlib.Path = None,
+    ) -> None:
+        if notebook_out is None:
+            notebook_out = notebook_in
+
+        notebook = nbformat.read(notebook_in, as_version=nbformat.NO_CONVERT)
+        notebook['metadata']['kernelspec']['name'] = kernel_name
+        nbformat.write(notebook, notebook_out)
+
+    def get_kernel_name(self, notebook_in: str | pathlib.Path) -> str:
+        notebook = nbformat.read(notebook_in, as_version=nbformat.NO_CONVERT)
+        return notebook['metadata']['kernelspec']['name']
+
+    def clear_outputs(
+        self, notebook_in: str | pathlib.Path, notebook_out: str | pathlib.Path = None
+    ) -> None:
+        if notebook_out is None:
+            notebook_out = notebook_in
+
+        notebook = nbformat.read(notebook_in, as_version=nbformat.NO_CONVERT)
+        if not isinstance(notebook['cells'], list):
+            raise TypeError('Notebook cells must be a list')
+
+        cells = []
+        for cell in notebook['cells']:
+            if cell['cell_type'] == 'code':
+                cell['execution_count'] = None
+                cell['outputs'] = []
+            cells.append(cell)
+        notebook['cells'] = cells
+        nbformat.write(notebook, notebook_out)
+
+    def execute(
+        self,
+        notebook_in: str | pathlib.Path,
+        output_dir: str | pathlib.Path = '.',
+        *,
+        timeout: int = None,
+    ) -> None | str:
+
+        with open(notebook_in, encoding='utf-8') as f:
+            notebook = nbformat.read(f, as_version=nbformat.NO_CONVERT)
+
+        executor = nbconvert.preprocessors.ExecutePreprocessor(timeout=timeout)
+
+        try:
+            output = executor.preprocess(notebook, {'metadata': {'path': './'}})
+        except nbconvert.preprocessors.CellExecutionError:
+            output = None
+            msg = f'Error executing the notebook "{notebook_in}".\n'
+            msg += f'See notebook "{notebook_in}" for the traceback.\n'
+            console.print(msg)
+
+        finally:
+            notebook_out = pathlib.Path(output_dir) / notebook_in.stem
+            with open(notebook_out, 'w', encoding='utf-8') as f:
+                nbformat.write(notebook, f)
+
+            console.print(f'Notebook "{notebook_out}" saved.')
+
+        return output