feat: migrate to awkward v2 (and keep supporting v1) (#284)

* chore: test on awkward v2.0.0rc1

* `RecordArrayType.recordlookup` -> `RecordArrayType.fields` for awkward v2

Co-authored-by: Angus Hollands <goosey15@gmail.com>

* Try splitting up CI jobs

* Fix doctests

* Better job names

Co-authored-by: Angus Hollands <goosey15@gmail.com>

* Even better job names

* Don't depend on pre-commit + --xdoctest

* Refactor looking up fields

Co-authored-by: Angus Hollands <goosey15@gmail.com>

* Use importlib_metadata only

* Remove python version restriction

* Better error handling

Co-authored-by: Angus Hollands <goosey15@gmail.com>

* Revert importlib.metdata changes

* Explicitly specify awkward v2

Co-authored-by: Angus Hollands <goosey15@gmail.com>

* Typo

Co-authored-by: Angus Hollands <goosey15@gmail.com>

* Typo

Co-authored-by: Angus Hollands <goosey15@gmail.com>

* Revert mypy config

* Try quotes

* Private function + uproot-awkward sync

* Changelog entry

* Try Python 3.8.15.

* I've confirmed that uncompyle6 loads for Python 3.8.13.

* Trigger the tests.

Co-authored-by: Angus Hollands <goosey15@gmail.com>
Co-authored-by: Jim Pivarski <jpivarski@gmail.com>

Author: 3 people

.github/CONTRIBUTING.md

@@ -122,6 +122,12 @@ The doctests can be executed using the `test` dependencies of `vector` in the fo
 xdoctest ./src/vector/
 ```
 
+or, one can run the doctests along with the unit tests in the following way -
+
+```bash
+python -m pytest --xdoctest .
+```
+
 A much more detailed guide on testing with `pytest` for `Scikit-HEP` packages is available [here](https://scikit-hep.org/developer/pytest).
 
 ### Running notebook tests

.github/PULL_REQUEST_TEMPLATE.md

@@ -7,12 +7,12 @@ _Please describe the purpose of this pull request. Reference and link to any rel
 ## Checklist
 
 - [ ] Have you followed the guidelines in our Contributing document?
-- [ ] Have you checked to ensure there aren't any other open Pull Request for the required change ?
+- [ ] Have you checked to ensure there aren't any other open Pull Requests for the required change?
 - [ ] Does your submission pass pre-commit? (`$ pre-commit run --all-files` or `$ nox -s lint`)
 - [ ] Does your submission pass tests? (`$ pytest` or `$ nox -s tests`)
 - [ ] Does the documentation build with your changes? (`$ cd docs; make clean; make html` or `$ nox -s docs`)
-- [ ] Does your submission pass doctests? (`$ xdoctest ./src/vector` or `$ nox -s doctests`)
+- [ ] Does your submission pass the doctests? (`$ xdoctest ./src/vector` or `$ nox -s doctests`)
 
 ## Before Merging
 
-- [ ] Summarize commit messages into a brief review of the Pull request.
+- [ ] Summarize the commit messages into a brief review of the Pull request.

.github/workflows/ci.yml

@@ -29,7 +29,7 @@ jobs:
           echo "::add-matcher::$GITHUB_WORKSPACE/.github/matchers/pylint.json"
           pipx run nox -s pylint
 
-  checks:
+  check-light:
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
@@ -40,7 +40,7 @@ jobs:
           - "3.9"
           - "3.10"
           - "3.11"
-    name: Check Python ${{ matrix.python-version }}
+    name: Python ${{ matrix.python-version }} - Light
     steps:
       - uses: actions/checkout@v3
 
@@ -57,20 +57,69 @@ jobs:
       - name: Test light package
         run: python -m pytest -ra --ignore tests/test_notebooks.py
 
-      - name: Install develop extras
+  check-awkward-v1:
+    needs: [check-light]
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version:
+          - "3.7"
+          - "3.8"
+          - "3.9"
+          - "3.10"
+          - "3.11"
+    name: Python ${{ matrix.python-version }} - Awkward v1
+    steps:
+      - uses: actions/checkout@v3
+
+      - uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Requirements check
+        run: python -m pip list
+
+      - name: Install package
         run: python -m pip install -e .[dev]
 
+      - name: Install awkward v1
+        run: python -m pip install -U "awkward<2"
+
       - name: Test package with awkward v1.x
-        run: python -m pytest -ra --cov=vector --ignore tests/test_notebooks.py tests/
+        run: python -m pytest -ra --cov=vector --xdoctest --ignore tests/test_notebooks.py .
 
-      - name: Use awkward v1.10.x for v2 support
-        run: python -m pip install -U --pre "awkward<2"
+  check-awkward-v2:
+    needs: [check-light]
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version:
+          - "3.7"
+          - "3.8"
+          - "3.9"
+          - "3.10"
+          - "3.11"
+    name: Python ${{ matrix.python-version }} - Awkward v2
+    steps:
+      - uses: actions/checkout@v3
+
+      - uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Requirements check
+        run: python -m pip list
+
+      - name: Install package
+        run: python -m pip install -e .[dev]
 
-      - name: Test package with awkward._v2
-        run: VECTOR_USE_AWKWARDV2=1 python -m pytest -ra --cov=vector --ignore tests/test_notebooks.py tests/
+      - name: Install awkward v2
+        run: python -m pip install -U --pre "awkward>=2.0.0rc1"
 
-      - name: Run doctests
-        run: xdoctest ./src/vector/
+      - name: Test package with awkward v2.x
+        run: python -m pytest -ra --cov=vector --xdoctest --ignore tests/test_notebooks.py .
 
       - name: Upload coverage report
         uses: codecov/codecov-action@v3.1.1
@@ -83,7 +132,7 @@ jobs:
 
       - uses: actions/setup-python@v4
         with:
-          python-version: 3.8.8
+          python-version: 3.8.13
 
       - name: Requirements check
         run: python -m pip list
@@ -95,7 +144,8 @@ jobs:
         run: python -m pytest -ra -m dis --ignore tests/test_notebooks.py
 
   pass:
-    needs: [pre-commit, checks, discheck]
+    needs:
+      [pre-commit, check-light, check-awkward-v1, check-awkward-v2, discheck]
     runs-on: ubuntu-latest
     steps:
       - run: echo "All jobs passed"

docs/changelog.md

@@ -23,7 +23,7 @@
 - Fix missing backslash in latex for readme [#285][]
 - chore: ignore flake8 B905 + improve bug report template [#297][]
 - chore: better and long term fix for flake8-bugbear [#298][]
-- feat: migrate to awkward v2 [#284][]
+- feat: migrate to awkward v2 (and keep supporting v1) [#284][]
 
 [#256]: https://github.com/scikit-hep/vector/pull/256
 [#254]: https://github.com/scikit-hep/vector/pull/254

environment.yml

@@ -10,7 +10,7 @@ dependencies:
   - numpy >=1.13.3
   - root >=6.18.04
   - pip:
-      - "awkward>=1.2.0"
-      - "uproot==4.*"
+      - "awkward>=2.0.0"
+      - "uproot>=5.0.0"
       - "scikit-hep-testdata>=0.2.0"
       - -e .

pyproject.toml

@@ -201,6 +201,6 @@ module = [
   "numba.*",
   "awkward.*",
 ]
+ignore_missing_imports = true
 disallow_untyped_defs = false
 disallow_untyped_calls = false
-ignore_missing_imports = true

src/vector/__init__.py

@@ -73,6 +73,13 @@ def _import_awkward() -> None:
         raise ImportError(msg)
 
 
+_is_awkward_v2: bool | None
+try:
+    _is_awkward_v2 = packaging.version.Version(
+        importlib_metadata.version("awkward")
+    ) >= packaging.version.Version("2.0.0rc1")
+except importlib_metadata.PackageNotFoundError:
+    _is_awkward_v2 = None
 try:
     import awkward
 

src/vector/backends/awkward.py

@@ -1598,6 +1598,16 @@ class MomentumRecord4D(MomentumAwkward4D, ak.Record):  # type: ignore[misc]
 
 # implementation of behaviors in Numba ########################################
 
+if vector._is_awkward_v2:
+
+    def _lookup_field(record_type: typing.Any, name: str) -> int:
+        return record_type.fields.index(name)
+
+else:
+
+    def _lookup_field(record_type: typing.Any, name: str) -> int:
+        return record_type.recordlookup.index(name)
+
 
 def _arraytype_of(awkwardtype: typing.Any, component: str) -> typing.Any:
     import numba
@@ -1633,36 +1643,36 @@ def _aztype_of(recordarraytype: typing.Any, is_momentum: bool) -> typing.Any:
 
     if is_momentum:
         try:
-            x_index = recordarraytype.recordlookup.index("px")
+            x_index = _lookup_field(recordarraytype, "px")
         except ValueError:
             x_index = None
     if x_index is None:
         try:
-            x_index = recordarraytype.recordlookup.index("x")
+            x_index = _lookup_field(recordarraytype, "x")
         except ValueError:
             x_index = None
     if is_momentum:
         try:
-            y_index = recordarraytype.recordlookup.index("py")
+            y_index = _lookup_field(recordarraytype, "py")
         except ValueError:
             y_index = None
     if y_index is None:
         try:
-            y_index = recordarraytype.recordlookup.index("y")
+            y_index = _lookup_field(recordarraytype, "y")
         except ValueError:
             y_index = None
     if is_momentum:
         try:
-            rho_index = recordarraytype.recordlookup.index("pt")
+            rho_index = _lookup_field(recordarraytype, "pt")
         except ValueError:
             rho_index = None
     if rho_index is None:
         try:
-            rho_index = recordarraytype.recordlookup.index("rho")
+            rho_index = _lookup_field(recordarraytype, "rho")
         except ValueError:
             rho_index = None
     try:
-        phi_index = recordarraytype.recordlookup.index("phi")
+        phi_index = _lookup_field(recordarraytype, "phi")
     except ValueError:
         phi_index = None
 
@@ -1702,20 +1712,20 @@ def _ltype_of(recordarraytype: typing.Any, is_momentum: bool) -> typing.Any:
 
     if is_momentum:
         try:
-            z_index = recordarraytype.recordlookup.index("pz")
+            z_index = _lookup_field(recordarraytype, "pz")
         except ValueError:
             z_index = None
     if z_index is None:
         try:
-            z_index = recordarraytype.recordlookup.index("z")
+            z_index = _lookup_field(recordarraytype, "z")
         except ValueError:
             z_index = None
     try:
-        theta_index = recordarraytype.recordlookup.index("theta")
+        theta_index = _lookup_field(recordarraytype, "theta")
     except ValueError:
         theta_index = None
     try:
-        eta_index = recordarraytype.recordlookup.index("eta")
+        eta_index = _lookup_field(recordarraytype, "eta")
     except ValueError:
         eta_index = None
 
@@ -1754,42 +1764,42 @@ def _ttype_of(recordarraytype: typing.Any, is_momentum: bool) -> typing.Any:
 
     if is_momentum:
         try:
-            t_index = recordarraytype.recordlookup.index("E")
+            t_index = _lookup_field(recordarraytype, "E")
         except ValueError:
             t_index = None
     if is_momentum and t_index is None:
         try:
-            t_index = recordarraytype.recordlookup.index("e")
+            t_index = _lookup_field(recordarraytype, "e")
         except ValueError:
             t_index = None
     if is_momentum and t_index is None:
         try:
-            t_index = recordarraytype.recordlookup.index("energy")
+            t_index = _lookup_field(recordarraytype, "energy")
         except ValueError:
             t_index = None
     if t_index is None:
         try:
-            t_index = recordarraytype.recordlookup.index("t")
+            t_index = _lookup_field(recordarraytype, "t")
         except ValueError:
             t_index = None
     if is_momentum:
         try:
-            tau_index = recordarraytype.recordlookup.index("M")
+            tau_index = _lookup_field(recordarraytype, "M")
         except ValueError:
             tau_index = None
     if is_momentum and tau_index is None:
         try:
-            tau_index = recordarraytype.recordlookup.index("m")
+            tau_index = _lookup_field(recordarraytype, "m")
         except ValueError:
             tau_index = None
     if is_momentum and tau_index is None:
         try:
-            tau_index = recordarraytype.recordlookup.index("mass")
+            tau_index = _lookup_field(recordarraytype, "mass")
         except ValueError:
             tau_index = None
     if tau_index is None:
         try:
-            tau_index = recordarraytype.recordlookup.index("tau")
+            tau_index = _lookup_field(recordarraytype, "tau")
         except ValueError:
             tau_index = None
 
@@ -1895,7 +1905,11 @@ def _numba_lower(
 
     vectorcls = sig.return_type.instance_class
 
-    fields = sig.args[0].arrayviewtype.type.recordlookup
+    fields = (
+        sig.args[0].arrayviewtype.type.fields
+        if vector._is_awkward_v2
+        else sig.args[0].arrayviewtype.type.recordlookup
+    )
 
     if issubclass(vectorcls, (VectorObject2D, VectorObject3D, VectorObject4D)):
         if issubclass(sig.return_type.azimuthaltype.instance_class, AzimuthalXY):

tests/backends/test_awkward.py

@@ -5,8 +5,6 @@
 
 from __future__ import annotations
 
-import os
-
 import pytest
 
 import vector
@@ -81,12 +79,6 @@ def test_rotateZ():
     assert out.wow.tolist() == [[99], [], [123]]
 
 
-# awkward._v2 has not yet registered NumPy dispatch mechanisms
-# see https://github.com/scikit-hep/awkward/issues/1638
-# TODO: ensure this passes once awkward v2 is out
-@pytest.mark.xfail(
-    strict=True if os.environ.get("VECTOR_USE_AWKWARDV2") is not None else False
-)
 def test_projection():
     array = vector.Array(
         [

tests/backends/test_awkward_numba.py

@@ -5,8 +5,6 @@
 
 from __future__ import annotations
 
-import os
-
 import pytest
 
 import vector
@@ -19,12 +17,6 @@
 pytestmark = [pytest.mark.numba, pytest.mark.awkward]
 
 
-# awkward._v2 has not yet registered Numba dispatch mechanisms
-# see https://github.com/scikit-hep/awkward/discussions/1639
-# TODO: ensure this passes once awkward v2 is out
-@pytest.mark.xfail(
-    strict=True if os.environ.get("VECTOR_USE_AWKWARDV2") is not None else False
-)
 def test():
     @numba.njit
     def extract(x):

tests/conftest.py

@@ -45,7 +45,6 @@
 spark_parser = pytest.importorskip("spark_parser")
 pytestmark = pytest.mark.dis
 
-
 Context = collections.namedtuple("Context", ["name", "closure"])