refactor: python doc script classes

.github/workflows/gh-ci-tests.yaml

@@ -43,7 +43,11 @@ jobs:
           micromamba info
           micromamba list
 
+      - name: "test snapshots"
+        working-directory: doc/source/scripts/
+        continue-on-error: true
+        run: python -m pytest
+
       - name: "test notebooks"
-        run: |
-          cd ${GITHUB_WORKSPACE}/tests
-          pytest
+        working-directory: tests/
+        run: python -m pytest

.pre-commit-config.yaml

@@ -5,7 +5,7 @@ repos:
     -   id: check-yaml
     -   id: end-of-file-fixer
     -   id: trailing-whitespace
-        exclude: ^.*\.(pdb)$
+        exclude: ^.*\.(pdb|ambr)$
 -   repo: https://github.com/pycqa/isort
     rev: 5.12.0
     hooks:
@@ -16,3 +16,9 @@ repos:
     rev: 22.3.0
     hooks:
     -   id: black
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.5.1
+    hooks:
+    -   id: mypy
+        args: [--ignore-missing-imports, --install-types, --non-interactive, --strict]
+        exclude: "/tests/.*\\.py|clean_example_notebooks.py|update_json_stubs_sitemap.py"

doc/source/conf.py

@@ -27,7 +27,7 @@
 project = "MDAnalysis User Guide"
 
 
-def sort_authors(filename):
+def sort_authors(filename: str) -> list[str]:
     """Generate sorted list of authors from AUTHORS"""
     authors = []
     with open(filename, "r") as f:
@@ -207,7 +207,7 @@ def sort_authors(filename):
 }
 
 # nbsphinx
-html_js_files = [
+html_js_files: list[str] = [
     # 'https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.4/require.min.js',
     # DEFAULT_EMBED_REQUIREJS_URL,
 ]

doc/source/scripts/base.py

@@ -1,15 +1,66 @@
-from __future__ import print_function
-
+import collections
 import os
 import pathlib
 import sys
 import textwrap
-from collections import defaultdict
+from collections.abc import Callable, Iterable
+from typing import Any, Optional, Type
 
+import pandas as pd
 import tabulate
 
 
-class TableWriter(object):
+def _run_method(method: Callable[..., str], *args: Any) -> str:
+    val = method(*args)
+    return val
+
+
+def _generate_row(
+    *, column_spec: list[tuple[str, Callable[..., str]]], args: Iterable[Any]
+) -> dict[str, str]:
+    row = {}
+    for heading, method in column_spec:
+        val = _run_method(method, *args)
+        row[heading] = val
+    return row
+
+
+def _generate_table(
+    *,
+    input_items: Iterable[Any],
+    column_spec: list[tuple[str, Callable[..., str]]],
+) -> pd.DataFrame:
+    rows = []
+    for args in input_items:
+        if not isinstance(args, Iterable):
+            args = [args]
+        line = _generate_row(column_spec=column_spec, args=args)
+        rows.append(line)
+    df = pd.DataFrame(rows)
+    return df
+
+
+def write_table(
+    *,
+    path: str,
+    headers: list[str],
+    lines: list[list[str]],
+    include_table: Optional[str] = None,
+) -> None:
+    parent_directory = pathlib.Path(path).parent
+    parent_directory.mkdir(exist_ok=True, parents=True)
+    with open(path, "w") as f:
+        f.write(f'..\n    Generated by {sys.argv[0].split("/")[-1]}\n\n')
+        if include_table:
+            f.write(f".. table:: {include_table}\n\n")
+        tabled = tabulate.tabulate(lines, headers=headers, tablefmt="rst")
+        if include_table:
+            tabled = textwrap.indent(tabled, "    ")
+        f.write(tabled)
+    print("Wrote ", path)
+
+
+class TableWriter:
     """
     For writing tables with easy column switching.
 
@@ -18,87 +69,73 @@ class TableWriter(object):
     Filename relative to source.
     """
 
-    filename = ""
-    include_table = False
-    headings = []
-    preprocess = []
-    postprocess = []
-    sort = True
+    def __init__(
+        self,
+        column_spec: list[tuple[str, Callable[..., str]]],
+        lines: list[list[str]],
+        filename: str = "",
+        include_table: Optional[str] = None,
+        sort: bool = True,
+        input_items: Optional[Iterable[Any]] = None,
+    ):
+        if column_spec:
+            assert input_items
+        assert (column_spec and not lines) or (lines and not column_spec)
+        stem = os.getcwd().split("source")[0]
+        self.path = os.path.join(stem, "source", filename)
+        self.filename = filename
+        self.include_table = include_table
+        self.sort = sort
+        self.input_items = input_items
+        self.column_spec = column_spec
+        self.lines = lines
+        self._df = pd.DataFrame()
 
-    def __getattr__(self, key: str) -> list:
-        return self.fields[key]
+    @property
+    def headers(self) -> list[str]:
+        return [column_name for column_name, _ in self.column_spec]
 
-    def __init__(self, *args, **kwargs):
-        stem = os.getcwd().split("source")[0]
-        self.path = os.path.join(stem, "source", self.filename)
-        self.fields = defaultdict(list)
+    @property
+    def fields(self) -> pd.DataFrame:
+        return self._df
 
-        parent_directory = pathlib.Path(self.path).parent
-        parent_directory.mkdir(exist_ok=True, parents=True)
-        self.get_lines(*args, **kwargs)
-        self.write_table()
+    def generate_lines_and_write_table(self) -> None:
+        df = _generate_table(
+            input_items=self.input_items or [],
+            column_spec=self.column_spec,
+        )
 
-    def _run_method(self, method, *args, **kwargs):
-        sanitized = self.sanitize_name(method)
-        meth = getattr(self, sanitized)
-        val = meth(*args, **kwargs)
-        self.fields[method].append(val)
-        return val
-
-    @staticmethod
-    def sanitize_name(name):
-        return "_" + name.replace(" ", "_").replace("/", "_").lower()
-
-    def get_lines(self, *args, **kwargs):
-        lines = []
-        for items in self._set_up_input():
-            try:
-                lines.append(self.get_line(*items))
-            except TypeError:  # one argument
-                lines.append(self.get_line(items))
-        if self.sort:
-            lines = sorted(lines)
+        lines = df.values.tolist()
+        lines = sorted(lines) if self.sort else lines
         self.lines = lines
+        self._df = df
+        self.write_table()
+
+    def write_table(self) -> None:
+        write_table(
+            path=self.path,
+            headers=self.headers,
+            lines=self.lines,
+            include_table=self.include_table,
+        )
+
+
+# ==== HELPER FUNCTIONS ==== #
+
+
+def sphinx_class(*, klass: Type[Any], tilde: bool = True) -> str:
+    prefix = "~" if tilde else ""
+    return f":class:`{prefix}{klass.__module__}.{klass.__name__}`"
+
+
+def sphinx_method(*, method: Callable[..., Any], tilde: bool = True) -> str:
+    prefix = "~" if tilde else ""
+    return ":meth:`{}{}.{}`".format(prefix, method.__module__, method.__qualname__)
+
+
+def sphinx_ref(*, txt: str, label: Optional[str] = None, suffix: str = "") -> str:
+    return f":ref:`{txt} <{label}{suffix}>`"
+
 
-    def get_line(self, *args):
-        line = []
-        for p in self.preprocess:
-            self._run_method(p, *args)
-        for h in self.headings:
-            line.append(self._run_method(h, *args))
-        for p in self.postprocess:
-            self._run_method(p, *args)
-        return line
-
-    def write_table(self):
-        with open(self.path, "w") as f:
-            f.write(f'..\n    Generated by {sys.argv[0].split("/")[-1]}\n\n')
-            if self.include_table:
-                f.write(f".. table:: {self.include_table}\n\n")
-            tabled = tabulate.tabulate(
-                self.lines, headers=self.headings, tablefmt="rst"
-            )
-            if self.include_table:
-                tabled = textwrap.indent(tabled, "    ")
-            f.write(tabled)
-        print("Wrote ", self.filename)
-
-    # ==== HELPER FUNCTIONS ==== #
-
-    @staticmethod
-    def sphinx_class(klass, tilde=True):
-        prefix = "~" if tilde else ""
-        return ":class:`{}{}.{}`".format(prefix, klass.__module__, klass.__name__)
-
-    @staticmethod
-    def sphinx_meth(meth, tilde=True):
-        prefix = "~" if tilde else ""
-        return ":meth:`{}{}.{}`".format(prefix, meth.__module__, meth.__qualname__)
-
-    @staticmethod
-    def sphinx_ref(txt: str, label: str = None, suffix: str = "") -> str:
-        return f":ref:`{txt} <{label}{suffix}>`"
-
-    @staticmethod
-    def sphinx_link(txt):
-        return "`{}`_".format(txt)
+def sphinx_link(*, txt: str) -> str:
+    return f"`{txt}`_"

doc/source/scripts/core.py

@@ -1,8 +1,3 @@
-from __future__ import print_function
-
-import os
-
-import tabulate
 from MDAnalysis import _TOPOLOGY_ATTRS
 
 # ====== TOPOLOGY ====== #