Skip to content

Commit 9071572

Browse files
Fix .pyi files after pybind11 v3 (#1106)
* Fix `.pyi` files after pybind11 v3 Fix the `impactx/impactx_pybind/__init__.pyi` file that references non-imported classes after the update to `pybind11` v3.0 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Patch `impactx/impactx_pybind/__init__.pyi` * Stabilize `distribution_input_helpers.py(i)` Move `as np` into body. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 1f3a721 commit 9071572

17 files changed

+143
-1573
lines changed

.github/update_stub.sh

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,7 @@ set -eu -o pipefail
1515
this_dir=$(cd $(dirname $0) && pwd)
1616

1717
pybind11-stubgen --exit-code -o ${this_dir}/../src/python/ impactx
18+
19+
# fix weird missing import issues after update to pybind11 v3.0
20+
sed -i 's/impactx.impactx_pybind.elements/elements/g' src/python/impactx/impactx_pybind/__init__.pyi
21+
sed -i 's/impactx.impactx_pybind.distribution/distribution/g' src/python/impactx/impactx_pybind/__init__.pyi

.pre-commit-config.yaml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,6 @@ repos:
7070
hooks:
7171
# Run the linter
7272
- id: ruff-check
73-
types_or: [ python ]
7473
args: [--fix, --exit-non-zero-on-fix]
7574
# Run the formatter
7675
- id: ruff-format

src/python/impactx/MADXParser.pyi

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from __future__ import annotations
2+
23
import os as os
34
import re as re
45
import warnings as warnings

src/python/impactx/__init__.pyi

Lines changed: 28 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -13,35 +13,42 @@ impactx_pybind
1313
"""
1414

1515
from __future__ import annotations
16+
17+
import os as os
18+
1619
from amrex import space3d as amr
1720
from amrex.space3d.amrex_3d_pybind import SmallMatrix_6x6_F_SI1_double as Map6x6
1821
from impactx.distribution_input_helpers import twiss
1922
from impactx.extensions.ImpactXParticleContainer import (
2023
register_ImpactXParticleContainer_extension,
2124
)
2225
from impactx.extensions.KnownElementsList import register_KnownElementsList_extension
23-
from impactx.impactx_pybind import Config
24-
from impactx.impactx_pybind import CoordSystem
25-
from impactx.impactx_pybind import Envelope
26-
from impactx.impactx_pybind import ImpactX
27-
from impactx.impactx_pybind import ImpactXParConstIter
28-
from impactx.impactx_pybind import ImpactXParIter
29-
from impactx.impactx_pybind import ImpactXParticleContainer
30-
from impactx.impactx_pybind import RefPart
31-
from impactx.impactx_pybind import coordinate_transformation
32-
from impactx.impactx_pybind import create_envelope
33-
from impactx.impactx_pybind import distribution
34-
from impactx.impactx_pybind import elements
35-
from impactx.impactx_pybind import push
36-
from impactx.impactx_pybind import wakeconvolution
26+
from impactx.impactx_pybind import (
27+
Config,
28+
CoordSystem,
29+
Envelope,
30+
ImpactX,
31+
ImpactXParConstIter,
32+
ImpactXParIter,
33+
ImpactXParticleContainer,
34+
RefPart,
35+
coordinate_transformation,
36+
create_envelope,
37+
distribution,
38+
elements,
39+
push,
40+
wakeconvolution,
41+
)
3742
from impactx.madx_to_impactx import read_beam
38-
import os as os
39-
from . import MADXParser
40-
from . import distribution_input_helpers
41-
from . import extensions
42-
from . import impactx_pybind
43-
from . import madx_to_impactx
44-
from . import plot
43+
44+
from . import (
45+
MADXParser,
46+
distribution_input_helpers,
47+
extensions,
48+
impactx_pybind,
49+
madx_to_impactx,
50+
plot,
51+
)
4552

4653
__all__: list[str] = [
4754
"Config",

src/python/impactx/distribution_input_helpers.py

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,19 +6,19 @@
66
#
77
# -*- coding: utf-8 -*-
88

9-
import numpy as np
9+
import numpy
1010

1111

1212
def twiss(
13-
beta_x: np.float64,
14-
beta_y: np.float64,
15-
beta_t: np.float64,
16-
emitt_x: np.float64,
17-
emitt_y: np.float64,
18-
emitt_t: np.float64,
19-
alpha_x: np.float64 = 0.0,
20-
alpha_y: np.float64 = 0.0,
21-
alpha_t: np.float64 = 0.0,
13+
beta_x: numpy.float64,
14+
beta_y: numpy.float64,
15+
beta_t: numpy.float64,
16+
emitt_x: numpy.float64,
17+
emitt_y: numpy.float64,
18+
emitt_t: numpy.float64,
19+
alpha_x: numpy.float64 = 0.0,
20+
alpha_y: numpy.float64 = 0.0,
21+
alpha_t: numpy.float64 = 0.0,
2222
):
2323
"""
2424
Helper function to convert Courant-Snyder / Twiss input into phase space ellipse input.
@@ -34,6 +34,8 @@ def twiss(
3434
:param alpha_t: Alpha function value in the t dimension, default is 0.0.
3535
:return: A dictionary containing calculated phase space input: 'lambdaX', 'lambdaY', 'lambdaT', 'lambdaPx', 'lambdaPy', 'lambdaPt', 'muxpx', 'muypy', 'mutpt'.
3636
"""
37+
import numpy as np
38+
3739
if beta_x <= 0.0 or beta_y <= 0.0 or beta_t <= 0.0:
3840
raise ValueError(
3941
"Input Error: The beta function values need to be non-zero positive values in all dimensions."

src/python/impactx/distribution_input_helpers.pyi

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
from __future__ import annotations
2-
import numpy as np
3-
import numpy
42

5-
__all__: list[str] = ["np", "twiss"]
3+
import numpy as numpy
4+
5+
__all__: list[str] = ["numpy", "twiss"]
66

77
def twiss(
88
beta_x: numpy.float64,

src/python/impactx/extensions/KnownElementsList.pyi

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,11 @@ License: BSD-3-Clause-LBNL
88
"""
99

1010
from __future__ import annotations
11-
from impactx.impactx_pybind import elements
11+
1212
import os as os
1313

14+
from impactx.impactx_pybind import elements
15+
1416
__all__: list[str] = [
1517
"elements",
1618
"from_pals",
Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from __future__ import annotations
2-
from . import ImpactXParticleContainer
3-
from . import KnownElementsList
2+
3+
from . import ImpactXParticleContainer, KnownElementsList
44

55
__all__: list[str] = ["ImpactXParticleContainer", "KnownElementsList"]

src/python/impactx/impactx_pybind/__init__.pyi

Lines changed: 70 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,15 @@ impactx_pybind
1313
"""
1414

1515
from __future__ import annotations
16-
from amrex import space3d as amr
17-
import amrex.space3d.amrex_3d_pybind
18-
from amrex.space3d.amrex_3d_pybind import SmallMatrix_6x6_F_SI1_double as Map6x6
16+
1917
import collections.abc
2018
import typing
21-
from . import distribution
22-
from . import elements
23-
from . import wakeconvolution
19+
20+
import amrex.space3d.amrex_3d_pybind
21+
from amrex import space3d as amr
22+
from amrex.space3d.amrex_3d_pybind import SmallMatrix_6x6_F_SI1_double as Map6x6
23+
24+
from . import distribution, elements, wakeconvolution
2425

2526
__all__: list[str] = [
2627
"Config",
@@ -108,15 +109,15 @@ class ImpactX:
108109
def add_particles(
109110
self,
110111
bunch_charge: typing.SupportsFloat,
111-
distr: impactx.impactx_pybind.distribution.Empty
112-
| impactx.impactx_pybind.distribution.Gaussian
113-
| impactx.impactx_pybind.distribution.Kurth4D
114-
| impactx.impactx_pybind.distribution.Kurth6D
115-
| impactx.impactx_pybind.distribution.KVdist
116-
| impactx.impactx_pybind.distribution.Thermal
117-
| impactx.impactx_pybind.distribution.Triangle
118-
| impactx.impactx_pybind.distribution.Semigaussian
119-
| impactx.impactx_pybind.distribution.Waterbag,
112+
distr: distribution.Empty
113+
| distribution.Gaussian
114+
| distribution.Kurth4D
115+
| distribution.Kurth6D
116+
| distribution.KVdist
117+
| distribution.Thermal
118+
| distribution.Triangle
119+
| distribution.Semigaussian
120+
| distribution.Waterbag,
120121
npart: typing.SupportsInt,
121122
) -> None:
122123
"""
@@ -145,15 +146,15 @@ class ImpactX:
145146
def init_envelope(
146147
self,
147148
ref: RefPart,
148-
distr: impactx.impactx_pybind.distribution.Empty
149-
| impactx.impactx_pybind.distribution.Gaussian
150-
| impactx.impactx_pybind.distribution.Kurth4D
151-
| impactx.impactx_pybind.distribution.Kurth6D
152-
| impactx.impactx_pybind.distribution.KVdist
153-
| impactx.impactx_pybind.distribution.Thermal
154-
| impactx.impactx_pybind.distribution.Triangle
155-
| impactx.impactx_pybind.distribution.Semigaussian
156-
| impactx.impactx_pybind.distribution.Waterbag,
149+
distr: distribution.Empty
150+
| distribution.Gaussian
151+
| distribution.Kurth4D
152+
| distribution.Kurth6D
153+
| distribution.KVdist
154+
| distribution.Thermal
155+
| distribution.Triangle
156+
| distribution.Semigaussian
157+
| distribution.Waterbag,
157158
intensity: typing.SupportsFloat | None = None,
158159
) -> None:
159160
"""
@@ -798,55 +799,55 @@ def coordinate_transformation(
798799
"""
799800

800801
def create_envelope(
801-
arg0: impactx.impactx_pybind.distribution.Empty
802-
| impactx.impactx_pybind.distribution.Gaussian
803-
| impactx.impactx_pybind.distribution.Kurth4D
804-
| impactx.impactx_pybind.distribution.Kurth6D
805-
| impactx.impactx_pybind.distribution.KVdist
806-
| impactx.impactx_pybind.distribution.Thermal
807-
| impactx.impactx_pybind.distribution.Triangle
808-
| impactx.impactx_pybind.distribution.Semigaussian
809-
| impactx.impactx_pybind.distribution.Waterbag,
802+
arg0: distribution.Empty
803+
| distribution.Gaussian
804+
| distribution.Kurth4D
805+
| distribution.Kurth6D
806+
| distribution.KVdist
807+
| distribution.Thermal
808+
| distribution.Triangle
809+
| distribution.Semigaussian
810+
| distribution.Waterbag,
810811
arg1: typing.SupportsFloat | None,
811812
) -> Envelope: ...
812813
def push(
813814
pc: ImpactXParticleContainer,
814-
element: impactx.impactx_pybind.elements.Empty
815-
| impactx.impactx_pybind.elements.Aperture
816-
| impactx.impactx_pybind.elements.Buncher
817-
| impactx.impactx_pybind.elements.CFbend
818-
| impactx.impactx_pybind.elements.ChrAcc
819-
| impactx.impactx_pybind.elements.ChrDrift
820-
| impactx.impactx_pybind.elements.ChrPlasmaLens
821-
| impactx.impactx_pybind.elements.ChrQuad
822-
| impactx.impactx_pybind.elements.ConstF
823-
| impactx.impactx_pybind.elements.BeamMonitor
824-
| impactx.impactx_pybind.elements.DipEdge
825-
| impactx.impactx_pybind.elements.Drift
826-
| impactx.impactx_pybind.elements.ExactCFbend
827-
| impactx.impactx_pybind.elements.ExactDrift
828-
| impactx.impactx_pybind.elements.ExactMultipole
829-
| impactx.impactx_pybind.elements.ExactQuad
830-
| impactx.impactx_pybind.elements.ExactSbend
831-
| impactx.impactx_pybind.elements.Kicker
832-
| impactx.impactx_pybind.elements.LinearMap
833-
| impactx.impactx_pybind.elements.Marker
834-
| impactx.impactx_pybind.elements.Multipole
835-
| impactx.impactx_pybind.elements.NonlinearLens
836-
| impactx.impactx_pybind.elements.PlaneXYRot
837-
| impactx.impactx_pybind.elements.Programmable
838-
| impactx.impactx_pybind.elements.PRot
839-
| impactx.impactx_pybind.elements.Quad
840-
| impactx.impactx_pybind.elements.QuadEdge
841-
| impactx.impactx_pybind.elements.RFCavity
842-
| impactx.impactx_pybind.elements.Sbend
843-
| impactx.impactx_pybind.elements.ShortRF
844-
| impactx.impactx_pybind.elements.SoftSolenoid
845-
| impactx.impactx_pybind.elements.SoftQuadrupole
846-
| impactx.impactx_pybind.elements.Sol
847-
| impactx.impactx_pybind.elements.Source
848-
| impactx.impactx_pybind.elements.TaperedPL
849-
| impactx.impactx_pybind.elements.ThinDipole,
815+
element: elements.Empty
816+
| elements.Aperture
817+
| elements.Buncher
818+
| elements.CFbend
819+
| elements.ChrAcc
820+
| elements.ChrDrift
821+
| elements.ChrPlasmaLens
822+
| elements.ChrQuad
823+
| elements.ConstF
824+
| elements.BeamMonitor
825+
| elements.DipEdge
826+
| elements.Drift
827+
| elements.ExactCFbend
828+
| elements.ExactDrift
829+
| elements.ExactMultipole
830+
| elements.ExactQuad
831+
| elements.ExactSbend
832+
| elements.Kicker
833+
| elements.LinearMap
834+
| elements.Marker
835+
| elements.Multipole
836+
| elements.NonlinearLens
837+
| elements.PlaneXYRot
838+
| elements.Programmable
839+
| elements.PRot
840+
| elements.Quad
841+
| elements.QuadEdge
842+
| elements.RFCavity
843+
| elements.Sbend
844+
| elements.ShortRF
845+
| elements.SoftSolenoid
846+
| elements.SoftQuadrupole
847+
| elements.Sol
848+
| elements.Source
849+
| elements.TaperedPL
850+
| elements.ThinDipole,
850851
step: typing.SupportsInt = 0,
851852
period: typing.SupportsInt = 0,
852853
) -> None:

src/python/impactx/impactx_pybind/distribution.pyi

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ Particle beam distributions in ImpactX
33
"""
44

55
from __future__ import annotations
6+
67
import typing
78

89
__all__: list[str] = [

0 commit comments

Comments
 (0)