feat: decorator to validate and annotate dependencies

docs/conf.py

@@ -42,6 +42,7 @@
 extensions = [
     "sphinx_copybutton",
     "sphinx_design",
+    "sphinx_togglebutton",
     "sphinx_external_toc",
     "sphinx.ext.intersphinx",
     "myst_nb",

docs/prepare_docstrings.py

@@ -2,13 +2,13 @@
 
 from __future__ import annotations
 
-import re
 import ast
 import glob
 import io
-import subprocess
-import pathlib
 import os
+import pathlib
+import re
+import subprocess
 import warnings
 
 import sphinx.ext.napoleon
@@ -20,7 +20,7 @@
 output_path.mkdir(exist_ok=True)
 
 latest_commit = (
-    subprocess.run(["git", "rev-parse", "HEAD"], stdout=subprocess.PIPE)
+    subprocess.run(["git", "rev-parse", "HEAD"], stdout=subprocess.PIPE, check=False)
     .stdout.decode("utf-8")
     .strip()
 )
@@ -48,16 +48,16 @@ def tostr(node):
         return tostr(node.value) + "." + node.attr
 
     elif isinstance(node, ast.Subscript):
-        return "{0}[{1}]".format(tostr(node.value), tostr(node.slice))
+        return f"{tostr(node.value)}[{tostr(node.slice)}]"
 
     elif isinstance(node, ast.Slice):
         start = "" if node.lower is None else tostr(node.lower)
         stop = "" if node.upper is None else tostr(node.upper)
         step = "" if node.step is None else tostr(node.step)
         if step == "":
-            return "{0}:{1}".format(start, stop)
+            return f"{start}:{stop}"
         else:
-            return "{0}:{1}:{2}".format(start, stop, step)
+            return f"{start}:{stop}:{step}"
 
     elif isinstance(node, ast.Call):
         return "{0}({1})".format(
@@ -78,10 +78,7 @@ def tostr(node):
 
     elif isinstance(node, ast.Dict):
         return "{{{0}}}".format(
-            ", ".join(
-                "{0}: {1}".format(tostr(x), tostr(y))
-                for x, y in zip(node.keys, node.values)
-            )
+            ", ".join(f"{tostr(x)}: {tostr(y)}" for x, y in zip(node.keys, node.values))
         )
 
     elif isinstance(node, ast.Lambda):
@@ -165,7 +162,7 @@ def dodoc(docstring, qualname, names):
     out = re.sub(r"<<<([^>]*)>>>", r"`\1`_", out)
     out = re.sub(r"#(ak\.[A-Za-z0-9_\.]*[A-Za-z0-9_])", r":py:obj:`\1`", out)
     for x in names:
-        out = out.replace("#" + x, ":py:meth:`{1} <{0}.{1}>`".format(qualname, x))
+        out = out.replace("#" + x, f":py:meth:`{x} <{qualname}.{x}>`")
     out = re.sub(r"\[([^\]]*)\]\(([^\)]*)\)", r"`\1 <\2>`__", out)
     out = str(sphinx.ext.napoleon.GoogleDocstring(out, config))
     out = re.sub(
@@ -212,25 +209,25 @@ def doclass(link, linelink, shortname, name, astcls):
     outfile = io.StringIO()
     outfile.write(qualname + "\n" + "-" * len(qualname) + "\n\n")
     outfile.write(f".. py:module: {qualname}\n\n")
-    outfile.write("Defined in {0}{1}.\n\n".format(link, linelink))
-    outfile.write(".. py:class:: {0}({1})\n\n".format(qualname, dosig(init)))
+    outfile.write(f"Defined in {link}{linelink}.\n\n")
+    outfile.write(f".. py:class:: {qualname}({dosig(init)})\n\n")
 
     docstring = ast.get_docstring(astcls)
     if docstring is not None:
         outfile.write(dodoc(docstring, qualname, names) + "\n\n")
 
     for node in rest:
         if isinstance(node, ast.Assign):
-            attrtext = "{0}.{1}".format(qualname, node.targets[0].id)
+            attrtext = f"{qualname}.{node.targets[0].id}"
             outfile.write(make_anchor(attrtext))
             outfile.write(".. py:attribute:: " + attrtext + "\n")
-            outfile.write("    :value: {0}\n\n".format(tostr(node.value)))
+            outfile.write(f"    :value: {tostr(node.value)}\n\n")
             docstring = None
 
         elif any(
             isinstance(x, ast.Name) and x.id == "property" for x in node.decorator_list
         ):
-            attrtext = "{0}.{1}".format(qualname, node.name)
+            attrtext = f"{qualname}.{node.name}"
             outfile.write(make_anchor(attrtext))
             outfile.write(".. py:attribute:: " + attrtext + "\n\n")
             docstring = ast.get_docstring(node)
@@ -242,8 +239,8 @@ def doclass(link, linelink, shortname, name, astcls):
             docstring = None
 
         else:
-            methodname = "{0}.{1}".format(qualname, node.name)
-            methodtext = "{0}({1})".format(methodname, dosig(node))
+            methodname = f"{qualname}.{node.name}"
+            methodtext = f"{methodname}({dosig(node)})"
             outfile.write(make_anchor(methodname))
             outfile.write(".. py:method:: " + methodtext + "\n\n")
             docstring = ast.get_docstring(node)
@@ -259,6 +256,56 @@ def doclass(link, linelink, shortname, name, astcls):
         entry_path.write_text(out)
 
 
+def get_function_dependency_spec(node: ast.FunctionDef):
+    for deco in node.decorator_list:
+        if (
+            isinstance(deco, ast.Call)
+            and (
+                (
+                    isinstance(deco.func, ast.Name)
+                    and deco.func.id == "high_level_function"
+                )
+                or (
+                    isinstance(deco.func, ast.Attribute)
+                    and isinstance(deco.func.value, ast.Name)
+                    and deco.func.value.id in ("ak", "awkward")
+                    and deco.func.attr == "high_level_function"
+                )
+            )
+            and (arg := {k.arg: k.value for k in deco.keywords}.get("dependencies"))
+            is not None
+        ):
+            break
+    else:
+        return None
+
+    extras = []
+    dependencies = []
+    if isinstance(arg, ast.Dict):
+        assert all(
+            isinstance(x, ast.Constant) and isinstance(x.value, str) for x in arg.keys
+        )
+        assert all(
+            isinstance(vals, (ast.List, ast.Tuple))
+            and all(
+                isinstance(x, ast.Constant) and isinstance(x.value, str)
+                for x in vals.elts
+            )
+            for vals in arg.values
+        )
+        extras.extend([x.value for x in arg.keys])
+        dependencies.extend([x.value for group in arg.values for x in group.elts])
+    elif isinstance(arg, (ast.List, ast.Tuple)):
+        assert all(
+            isinstance(x, ast.Constant) and isinstance(x.value, str) for x in arg.elts
+        )
+        dependencies.extend([x.value for x in arg.elts])
+    else:
+        raise TypeError(arg)
+
+    return extras, dependencies
+
+
 def dofunction(link, linelink, shortname, name, astfcn):
     if name.startswith("_"):
         return
@@ -268,15 +315,46 @@ def dofunction(link, linelink, shortname, name, astfcn):
     outfile = io.StringIO()
     outfile.write(qualname + "\n" + "-" * len(qualname) + "\n\n")
     outfile.write(f".. py:module: {qualname}\n\n")
-    outfile.write("Defined in {0}{1}.\n\n".format(link, linelink))
+    outfile.write(f"Defined in {link}{linelink}.\n\n")
 
-    functiontext = "{0}({1})".format(qualname, dosig(astfcn))
+    functiontext = f"{qualname}({dosig(astfcn)})"
     outfile.write(".. py:function:: " + functiontext + "\n\n")
 
     docstring = ast.get_docstring(astfcn)
     if docstring is not None:
         outfile.write(dodoc(docstring, qualname, []) + "\n\n")
 
+    dependency_spec = get_function_dependency_spec(astfcn)
+    if dependency_spec is not None:
+        extras, dependencies = dependency_spec
+
+        outfile.write(
+            ".. note::\n"
+            "   :class: dropdown\n"
+            "\n"
+            "   This function requires the following optional third-party libraries:\n"
+            "\n"
+        )
+        outfile.write("\n".join([f"   * ``{dep}``" for dep in dependencies]))
+        install_dependencies_string = " ".join(dependencies)
+        outfile.write(
+            "\n\n"
+            f"   If you use pip, you can install these packages with "
+            f"``python -m pip install {install_dependencies_string}``.\n"
+            "   Otherwise, if you use Conda, install the corresponding packages "
+            "for the correct versions. "
+        )
+        if extras:
+            extras_string = ",".join(extras)
+            outfile.write(
+                "\n\n   These dependencies can also be conveniently installed using the following extras:\n\n"
+            )
+            outfile.write("\n".join([f"   * ``{extra}``" for extra in extras]))
+            outfile.write(
+                "\n\n   If you use ``pip`` to install your packages, these extras can be installed by running "
+                f"``python -m pip install awkward[{extras_string}]``."
+            )
+
     out = outfile.getvalue()
 
     toctree.append(os.path.join("generated", qualname + ".rst"))
@@ -384,7 +462,6 @@ def func(x, *, y, z=None):
 test_signature_vararg()
 test_signature_kwonly()
 
-
 toctree_path = reference_path / "toctree.txt"
 toctree_contents = toctree_path.read_text()
 

docs/requirements.txt

@@ -8,9 +8,19 @@ PyYAML
 black
 pycparser
 lark-parser
+
+# Pin this
+sphinxcontrib-applehelp==1.0.0
+sphinxcontrib-devhelp==1.0.0
+sphinxcontrib-htmlhelp==2.0.0
+sphinxcontrib-jsmath==1.0.1
+sphinxcontrib-qthelp==1.0.0
+sphinxcontrib-serializinghtml==1.1.5
+
 sphinx-copybutton
 sphinx-design
 sphinx-sitemap
+sphinx-togglebutton
 pydata-sphinx-theme
 myst-nb
 sphinx-external-toc

pyproject.toml

@@ -50,6 +50,14 @@ dynamic = [
     "readme"
 ]
 
+[project.optional-dependencies]
+arrow = [
+    "arrow>=7.0.0",
+    "fsspec"
+]
+
+
+
 [project.entry-points.numba_extensions]
 init = "awkward.numba:_register"
 

requirements-test-full.txt

@@ -1,4 +1,4 @@
-fsspec;sys_platform != "win32"
+fsspec
 jax[cpu]>=0.2.15;sys_platform != "win32" and python_version < "3.12"
 numba>=0.50.0,!=0.58.0rc1;python_version < "3.12"
 numexpr>=2.7; python_version < "3.12"

requirements-test-minimal.txt

@@ -1,4 +1,4 @@
-fsspec;sys_platform != "win32"
+fsspec
 numpy==1.18.0
 pyarrow==7.0.0
 pytest>=6

src/awkward/_connect/pyarrow.py

@@ -14,75 +14,35 @@
 from awkward._nplikes.numpy import Numpy
 from awkward._nplikes.numpy_like import NumpyMetadata
 from awkward._parameters import parameters_union
+from awkward._requirements import import_required_module
 
 np = NumpyMetadata.instance()
 numpy = Numpy.instance()
 
 try:
     import pyarrow
 
-    error_message = None
+    if parse_version(pyarrow.__version__) < parse_version("7.0.0"):
+        raise ImportError
 
-except ModuleNotFoundError:
+except (ModuleNotFoundError, ImportError):
     pyarrow = None
-    error_message = """to use {0}, you must install pyarrow:
-
-    pip install pyarrow
-
-or
-
-    conda install -c conda-forge pyarrow
-"""
-
-else:
-    if parse_version(pyarrow.__version__) < parse_version("7.0.0"):
-        pyarrow = None
-        error_message = "pyarrow 7.0.0 or later required for {0}"
 
 
 def import_pyarrow(name: str) -> ModuleType:
-    if pyarrow is None:
-        raise ImportError(error_message.format(name))
-    return pyarrow
+    return import_required_module("pyarrow")
 
 
 def import_pyarrow_parquet(name: str) -> ModuleType:
-    if pyarrow is None:
-        raise ImportError(error_message.format(name))
-
-    import pyarrow.parquet as out
-
-    return out
+    return import_required_module("pyarrow.parquet")
 
 
 def import_pyarrow_compute(name: str) -> ModuleType:
-    if pyarrow is None:
-        raise ImportError(error_message.format(name))
-
-    import pyarrow.compute as out
-
-    return out
+    return import_required_module("pyarrow.compute")
 
 
 def import_fsspec(name: str) -> ModuleType:
-    try:
-        import fsspec
-
-    except ModuleNotFoundError as err:
-        raise ImportError(
-            f"""to use {name}, you must install fsspec:
-
-    pip install fsspec
-
-or
-
-    conda install -c conda-forge fsspec
-"""
-        ) from err
-
-    import_pyarrow_parquet(name)
-
-    return fsspec
+    return import_required_module("fsspec")
 
 
 if pyarrow is not None: