atlas-schema 0.2.1__tar.gz → 0.2.2__tar.gz

{atlas_schema-0.2.1 → atlas_schema-0.2.2}/PKG-INFO

@@ -1,11 +1,13 @@
 Metadata-Version: 2.4
 Name: atlas-schema
-Version: 0.2.1
+Version: 0.2.2
 Summary: Helper python package for ATLAS Common NTuple Analysis work.
 Project-URL: Homepage, https://github.com/scipp-atlas/atlas-schema
 Project-URL: Bug Tracker, https://github.com/scipp-atlas/atlas-schema/issues
 Project-URL: Discussions, https://github.com/scipp-atlas/atlas-schema/discussions
-Project-URL: Changelog, https://github.com/scipp-atlas/atlas-schema/releases
+Project-URL: Documentation, https://atlas-schema.readthedocs.io/en/v0.2.2/
+Project-URL: Releases, https://github.com/scipp-atlas/atlas-schema/releases
+Project-URL: Release Notes, https://atlas-schema.readthedocs.io/en/latest/history.html
 Author-email: Giordon Stark <kratsg@gmail.com>
 License: 
                                          Apache License
@@ -249,7 +251,7 @@ Requires-Dist: tbump>=6.7.0; extra == 'test'
 Requires-Dist: twine; extra == 'test'
 Description-Content-Type: text/markdown
 
-# atlas-schema v0.2.1
+# atlas-schema v0.2.2
 
 [![Actions Status][actions-badge]][actions-link]
 [![Documentation Status][rtd-badge]][rtd-link]

{atlas_schema-0.2.1 → atlas_schema-0.2.2}/README.md

@@ -1,4 +1,4 @@
-# atlas-schema v0.2.1
+# atlas-schema v0.2.2
 
 [![Actions Status][actions-badge]][actions-link]
 [![Documentation Status][rtd-badge]][rtd-link]

{atlas_schema-0.2.1 → atlas_schema-0.2.2}/pyproject.toml

@@ -60,7 +60,9 @@ docs = [
 Homepage = "https://github.com/scipp-atlas/atlas-schema"
 "Bug Tracker" = "https://github.com/scipp-atlas/atlas-schema/issues"
 Discussions = "https://github.com/scipp-atlas/atlas-schema/discussions"
-Changelog = "https://github.com/scipp-atlas/atlas-schema/releases"
+Documentation = "https://atlas-schema.readthedocs.io/en/v0.2.2/"
+Releases = "https://github.com/scipp-atlas/atlas-schema/releases"
+"Release Notes" = "https://atlas-schema.readthedocs.io/en/latest/history.html"
 
 
 [tool.hatch]
@@ -94,7 +96,19 @@ packages = ["src/atlas_schema"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"]
+addopts = [
+  "-ra",
+  "--showlocals",
+  "--strict-markers",
+  "--strict-config",
+  "--doctest-modules",
+  "--doctest-glob=*.rst",
+  "--cov",
+  "--cov-report=xml",
+  "--cov-report=term",
+  "--durations=20",
+  "--ignore=docs/conf.py",
+]
 xfail_strict = true
 filterwarnings = [
   "error",
@@ -102,9 +116,9 @@ filterwarnings = [
 log_cli_level = "INFO"
 testpaths = [
   "tests",
+  "docs",
 ]
 
-
 [tool.coverage]
 run.source = ["atlas_schema"]
 report.exclude_also = [
@@ -113,6 +127,7 @@ report.exclude_also = [
 ]
 
 [tool.mypy]
+mypy_path = ["src"]
 files = ["src", "tests"]
 python_version = "3.9"
 warn_unused_configs = true
@@ -133,6 +148,7 @@ disallow_incomplete_defs = true
 module = [
   'awkward.*',
   'coffea.*',
+  'dask_awkward.*',
 ]
 ignore_missing_imports = true
 

{atlas_schema-0.2.1 → atlas_schema-0.2.2}/src/atlas_schema/__init__.py

@@ -10,7 +10,8 @@ import warnings
 
 from atlas_schema._version import version as __version__
 from atlas_schema.enums import ParticleOrigin, PhotonID
+from atlas_schema.utils import isin
 
 warnings.filterwarnings("ignore", module="coffea.*")
 
-__all__ = ["ParticleOrigin", "PhotonID", "__version__"]
+__all__ = ["ParticleOrigin", "PhotonID", "__version__", "isin"]

{atlas_schema-0.2.1 → atlas_schema-0.2.2}/src/atlas_schema/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1'
-__version_tuple__ = version_tuple = (0, 2, 1)
+__version__ = version = '0.2.2'
+__version_tuple__ = version_tuple = (0, 2, 2)

{atlas_schema-0.2.1 → atlas_schema-0.2.2}/src/atlas_schema/enums.py

@@ -1,9 +1,31 @@
 from __future__ import annotations
 
-from enum import IntEnum
+import sys
+from enum import Enum, IntEnum
 
+if sys.version_info >= (3, 11):
+    from enum import EnumType
+else:
+    from enum import EnumMeta as EnumType
 
-class ParticleType(IntEnum):
+from typing import Callable, TypeVar, cast
+
+_E = TypeVar("_E", bound=Enum)
+
+
+class MultipleEnumAccessMeta(EnumType):
+    """
+    Enum Metaclass to provide a way to access multiple values all at once.
+    """
+
+    def __getitem__(self: type[_E], key: str | tuple[str]) -> _E | list[_E]:  # type:ignore[misc,override]
+        getitem = cast(Callable[[str], _E], super().__getitem__)  # type:ignore[misc]
+        if isinstance(key, tuple):
+            return [getitem(name) for name in key]
+        return getitem(key)
+
+
+class ParticleType(IntEnum, metaclass=MultipleEnumAccessMeta):
     """
     Taken from `ATLAS Truth Utilities for ParticleType <https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L8-49>`_.
     """
@@ -50,7 +72,7 @@ class ParticleType(IntEnum):
     UnknownJet = 38
 
 
-class ParticleOrigin(IntEnum):
+class ParticleOrigin(IntEnum, metaclass=MultipleEnumAccessMeta):
     """
     Taken from `ATLAS Truth Utilities for ParticleOrigin <https://gitlab.cern.ch/atlas/athena/-/blob/74f43ff0910edb2a2bd3778880ccbdad648dc037/Generators/TruthUtils/TruthUtils/TruthClasses.h#L51-103>`_.
     """
@@ -105,7 +127,7 @@ class ParticleOrigin(IntEnum):
     QCD = 45
 
 
-class PhotonID(IntEnum):
+class PhotonID(IntEnum, metaclass=MultipleEnumAccessMeta):
     """
     Taken from the `EGamma Identification CP group's twiki <https://twiki.cern.ch/twiki/bin/viewauth/AtlasProtected/EGammaIdentificationRun2#Photon_isEM_word>`_.
     """

atlas_schema-0.2.2/src/atlas_schema/utils.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import TypeVar, Union, cast
+
+import awkward as ak
+import dask_awkward as dak
+
+Array = TypeVar("Array", bound=Union[dak.Array, ak.Array])
+_E = TypeVar("_E", bound=Enum)
+
+
+def isin(haystack: Array, needles: dak.Array | ak.Array, axis: int = -1) -> Array:
+    """
+    Find needles in haystack.
+
+    This works by first transforming needles to an array with one more
+    dimension than the haystack, placing the needles at axis, and then doing a
+    comparison.
+
+    Args:
+        haystack (dak.Array or ak.Array): haystack of values.
+        needles (dak.Array or ak.Array): one-dimensional set of needles to find in haystack.
+        axis (int): the axis along which the comparison is performed
+
+    Returns:
+        dak.Array or ak.Array: result of comparison for needles in haystack
+    """
+    assert needles.ndim == 1, "Needles must be one-dimensional"
+    assert axis >= -1, "axis must be -1 or positive-valued"
+    assert axis < haystack.ndim + 1, "axis too large for the haystack"
+
+    # First, build up the transformation, with slice(None) indicating where to stick the needles
+    reshaper: list[None | slice] = [None] * haystack.ndim
+    axis = haystack.ndim if axis == -1 else axis
+    reshaper.insert(axis, slice(None))
+
+    # Note: reshaper needs to be a tuple for indexing purposes
+    return cast(Array, ak.any(haystack == needles[tuple(reshaper)], axis=-1))