Fix .pyi files after pybind11 v3

.github/update_stub.sh

@@ -15,3 +15,7 @@ set -eu -o pipefail
 this_dir=$(cd $(dirname $0) && pwd)
 
 pybind11-stubgen --exit-code -o ${this_dir}/../src/python/ impactx
+
+# fix weird missing import issues after update to pybind11 v3.0
+sed -i 's/impactx.impactx_pybind.elements/elements/g' src/python/impactx/impactx_pybind/__init__.pyi
+sed -i 's/impactx.impactx_pybind.distribution/distribution/g' src/python/impactx/impactx_pybind/__init__.pyi

.pre-commit-config.yaml

@@ -70,7 +70,6 @@ repos:
   hooks:
     # Run the linter
     - id: ruff-check
-      types_or: [ python ]
       args: [--fix, --exit-non-zero-on-fix]
     # Run the formatter
     - id: ruff-format

src/python/impactx/MADXParser.pyi

@@ -1,4 +1,5 @@
 from __future__ import annotations
+
 import os as os
 import re as re
 import warnings as warnings

src/python/impactx/__init__.pyi

@@ -13,35 +13,42 @@ impactx_pybind
 """
 
 from __future__ import annotations
+
+import os as os
+
 from amrex import space3d as amr
 from amrex.space3d.amrex_3d_pybind import SmallMatrix_6x6_F_SI1_double as Map6x6
 from impactx.distribution_input_helpers import twiss
 from impactx.extensions.ImpactXParticleContainer import (
     register_ImpactXParticleContainer_extension,
 )
 from impactx.extensions.KnownElementsList import register_KnownElementsList_extension
-from impactx.impactx_pybind import Config
-from impactx.impactx_pybind import CoordSystem
-from impactx.impactx_pybind import Envelope
-from impactx.impactx_pybind import ImpactX
-from impactx.impactx_pybind import ImpactXParConstIter
-from impactx.impactx_pybind import ImpactXParIter
-from impactx.impactx_pybind import ImpactXParticleContainer
-from impactx.impactx_pybind import RefPart
-from impactx.impactx_pybind import coordinate_transformation
-from impactx.impactx_pybind import create_envelope
-from impactx.impactx_pybind import distribution
-from impactx.impactx_pybind import elements
-from impactx.impactx_pybind import push
-from impactx.impactx_pybind import wakeconvolution
+from impactx.impactx_pybind import (
+    Config,
+    CoordSystem,
+    Envelope,
+    ImpactX,
+    ImpactXParConstIter,
+    ImpactXParIter,
+    ImpactXParticleContainer,
+    RefPart,
+    coordinate_transformation,
+    create_envelope,
+    distribution,
+    elements,
+    push,
+    wakeconvolution,
+)
 from impactx.madx_to_impactx import read_beam
-import os as os
-from . import MADXParser
-from . import distribution_input_helpers
-from . import extensions
-from . import impactx_pybind
-from . import madx_to_impactx
-from . import plot
+
+from . import (
+    MADXParser,
+    distribution_input_helpers,
+    extensions,
+    impactx_pybind,
+    madx_to_impactx,
+    plot,
+)
 
 __all__: list[str] = [
     "Config",

src/python/impactx/distribution_input_helpers.py

@@ -6,19 +6,19 @@
 #
 # -*- coding: utf-8 -*-
 
-import numpy as np
+import numpy
 
 
 def twiss(
-    beta_x: np.float64,
-    beta_y: np.float64,
-    beta_t: np.float64,
-    emitt_x: np.float64,
-    emitt_y: np.float64,
-    emitt_t: np.float64,
-    alpha_x: np.float64 = 0.0,
-    alpha_y: np.float64 = 0.0,
-    alpha_t: np.float64 = 0.0,
+    beta_x: numpy.float64,
+    beta_y: numpy.float64,
+    beta_t: numpy.float64,
+    emitt_x: numpy.float64,
+    emitt_y: numpy.float64,
+    emitt_t: numpy.float64,
+    alpha_x: numpy.float64 = 0.0,
+    alpha_y: numpy.float64 = 0.0,
+    alpha_t: numpy.float64 = 0.0,
 ):
     """
     Helper function to convert Courant-Snyder / Twiss input into phase space ellipse input.
@@ -34,6 +34,8 @@ def twiss(
     :param alpha_t: Alpha function value in the t dimension, default is 0.0.
     :return: A dictionary containing calculated phase space input: 'lambdaX', 'lambdaY', 'lambdaT', 'lambdaPx', 'lambdaPy', 'lambdaPt', 'muxpx', 'muypy', 'mutpt'.
     """
+    import numpy as np
+
     if beta_x <= 0.0 or beta_y <= 0.0 or beta_t <= 0.0:
         raise ValueError(
             "Input Error: The beta function values need to be non-zero positive values in all dimensions."

src/python/impactx/distribution_input_helpers.pyi

@@ -1,8 +1,8 @@
 from __future__ import annotations
-import numpy as np
-import numpy
 
-__all__: list[str] = ["np", "twiss"]
+import numpy as numpy
+
+__all__: list[str] = ["numpy", "twiss"]
 
 def twiss(
     beta_x: numpy.float64,

src/python/impactx/extensions/KnownElementsList.pyi

@@ -8,9 +8,11 @@ License: BSD-3-Clause-LBNL
 """
 
 from __future__ import annotations
-from impactx.impactx_pybind import elements
+
 import os as os
 
+from impactx.impactx_pybind import elements
+
 __all__: list[str] = [
     "elements",
     "from_pals",

src/python/impactx/extensions/__init__.pyi

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from . import ImpactXParticleContainer
-from . import KnownElementsList
+
+from . import ImpactXParticleContainer, KnownElementsList
 
 __all__: list[str] = ["ImpactXParticleContainer", "KnownElementsList"]

src/python/impactx/impactx_pybind/__init__.pyi

@@ -13,14 +13,15 @@ impactx_pybind
 """
 
 from __future__ import annotations
-from amrex import space3d as amr
-import amrex.space3d.amrex_3d_pybind
-from amrex.space3d.amrex_3d_pybind import SmallMatrix_6x6_F_SI1_double as Map6x6
+
 import collections.abc
 import typing
-from . import distribution
-from . import elements
-from . import wakeconvolution
+
+import amrex.space3d.amrex_3d_pybind
+from amrex import space3d as amr
+from amrex.space3d.amrex_3d_pybind import SmallMatrix_6x6_F_SI1_double as Map6x6
+
+from . import distribution, elements, wakeconvolution
 
 __all__: list[str] = [
     "Config",
@@ -108,15 +109,15 @@ class ImpactX:
     def add_particles(
         self,
         bunch_charge: typing.SupportsFloat,
-        distr: impactx.impactx_pybind.distribution.Empty
-        | impactx.impactx_pybind.distribution.Gaussian
-        | impactx.impactx_pybind.distribution.Kurth4D
-        | impactx.impactx_pybind.distribution.Kurth6D
-        | impactx.impactx_pybind.distribution.KVdist
-        | impactx.impactx_pybind.distribution.Thermal
-        | impactx.impactx_pybind.distribution.Triangle
-        | impactx.impactx_pybind.distribution.Semigaussian
-        | impactx.impactx_pybind.distribution.Waterbag,
+        distr: distribution.Empty
+        | distribution.Gaussian
+        | distribution.Kurth4D
+        | distribution.Kurth6D
+        | distribution.KVdist
+        | distribution.Thermal
+        | distribution.Triangle
+        | distribution.Semigaussian
+        | distribution.Waterbag,
         npart: typing.SupportsInt,
     ) -> None:
         """
@@ -145,15 +146,15 @@ class ImpactX:
     def init_envelope(
         self,
         ref: RefPart,
-        distr: impactx.impactx_pybind.distribution.Empty
-        | impactx.impactx_pybind.distribution.Gaussian
-        | impactx.impactx_pybind.distribution.Kurth4D
-        | impactx.impactx_pybind.distribution.Kurth6D
-        | impactx.impactx_pybind.distribution.KVdist
-        | impactx.impactx_pybind.distribution.Thermal
-        | impactx.impactx_pybind.distribution.Triangle
-        | impactx.impactx_pybind.distribution.Semigaussian
-        | impactx.impactx_pybind.distribution.Waterbag,
+        distr: distribution.Empty
+        | distribution.Gaussian
+        | distribution.Kurth4D
+        | distribution.Kurth6D
+        | distribution.KVdist
+        | distribution.Thermal
+        | distribution.Triangle
+        | distribution.Semigaussian
+        | distribution.Waterbag,
         intensity: typing.SupportsFloat | None = None,
     ) -> None:
         """
@@ -798,55 +799,55 @@ def coordinate_transformation(
     """
 
 def create_envelope(
-    arg0: impactx.impactx_pybind.distribution.Empty
-    | impactx.impactx_pybind.distribution.Gaussian
-    | impactx.impactx_pybind.distribution.Kurth4D
-    | impactx.impactx_pybind.distribution.Kurth6D
-    | impactx.impactx_pybind.distribution.KVdist
-    | impactx.impactx_pybind.distribution.Thermal
-    | impactx.impactx_pybind.distribution.Triangle
-    | impactx.impactx_pybind.distribution.Semigaussian
-    | impactx.impactx_pybind.distribution.Waterbag,
+    arg0: distribution.Empty
+    | distribution.Gaussian
+    | distribution.Kurth4D
+    | distribution.Kurth6D
+    | distribution.KVdist
+    | distribution.Thermal
+    | distribution.Triangle
+    | distribution.Semigaussian
+    | distribution.Waterbag,
     arg1: typing.SupportsFloat | None,
 ) -> Envelope: ...
 def push(
     pc: ImpactXParticleContainer,
-    element: impactx.impactx_pybind.elements.Empty
-    | impactx.impactx_pybind.elements.Aperture
-    | impactx.impactx_pybind.elements.Buncher
-    | impactx.impactx_pybind.elements.CFbend
-    | impactx.impactx_pybind.elements.ChrAcc
-    | impactx.impactx_pybind.elements.ChrDrift
-    | impactx.impactx_pybind.elements.ChrPlasmaLens
-    | impactx.impactx_pybind.elements.ChrQuad
-    | impactx.impactx_pybind.elements.ConstF
-    | impactx.impactx_pybind.elements.BeamMonitor
-    | impactx.impactx_pybind.elements.DipEdge
-    | impactx.impactx_pybind.elements.Drift
-    | impactx.impactx_pybind.elements.ExactCFbend
-    | impactx.impactx_pybind.elements.ExactDrift
-    | impactx.impactx_pybind.elements.ExactMultipole
-    | impactx.impactx_pybind.elements.ExactQuad
-    | impactx.impactx_pybind.elements.ExactSbend
-    | impactx.impactx_pybind.elements.Kicker
-    | impactx.impactx_pybind.elements.LinearMap
-    | impactx.impactx_pybind.elements.Marker
-    | impactx.impactx_pybind.elements.Multipole
-    | impactx.impactx_pybind.elements.NonlinearLens
-    | impactx.impactx_pybind.elements.PlaneXYRot
-    | impactx.impactx_pybind.elements.Programmable
-    | impactx.impactx_pybind.elements.PRot
-    | impactx.impactx_pybind.elements.Quad
-    | impactx.impactx_pybind.elements.QuadEdge
-    | impactx.impactx_pybind.elements.RFCavity
-    | impactx.impactx_pybind.elements.Sbend
-    | impactx.impactx_pybind.elements.ShortRF
-    | impactx.impactx_pybind.elements.SoftSolenoid
-    | impactx.impactx_pybind.elements.SoftQuadrupole
-    | impactx.impactx_pybind.elements.Sol
-    | impactx.impactx_pybind.elements.Source
-    | impactx.impactx_pybind.elements.TaperedPL
-    | impactx.impactx_pybind.elements.ThinDipole,
+    element: elements.Empty
+    | elements.Aperture
+    | elements.Buncher
+    | elements.CFbend
+    | elements.ChrAcc
+    | elements.ChrDrift
+    | elements.ChrPlasmaLens
+    | elements.ChrQuad
+    | elements.ConstF
+    | elements.BeamMonitor
+    | elements.DipEdge
+    | elements.Drift
+    | elements.ExactCFbend
+    | elements.ExactDrift
+    | elements.ExactMultipole
+    | elements.ExactQuad
+    | elements.ExactSbend
+    | elements.Kicker
+    | elements.LinearMap
+    | elements.Marker
+    | elements.Multipole
+    | elements.NonlinearLens
+    | elements.PlaneXYRot
+    | elements.Programmable
+    | elements.PRot
+    | elements.Quad
+    | elements.QuadEdge
+    | elements.RFCavity
+    | elements.Sbend
+    | elements.ShortRF
+    | elements.SoftSolenoid
+    | elements.SoftQuadrupole
+    | elements.Sol
+    | elements.Source
+    | elements.TaperedPL
+    | elements.ThinDipole,
     step: typing.SupportsInt = 0,
     period: typing.SupportsInt = 0,
 ) -> None:

src/python/impactx/impactx_pybind/distribution.pyi

@@ -3,6 +3,7 @@ Particle beam distributions in ImpactX
 """
 
 from __future__ import annotations
+
 import typing
 
 __all__: list[str] = [