From eb8592e4ec5fe294f27722ec8426461d97443c29 Mon Sep 17 00:00:00 2001 From: Mateusz Jakub Fila <37295697+m-fila@users.noreply.github.com> Date: Mon, 30 Sep 2024 11:07:21 +0200 Subject: [PATCH] add formatting and linting python in pre-commit (#370) * add ruff pre-commit * install ruff in pre-commit venv * add cmake target for ruff and ruff-format * fix formatting * fix linted warnings * be explicit about formatting options * target py311 --- .github/workflows/pre-commit.yml | 2 +- .pre-commit-config.yaml | 10 ++++ .ruff.toml | 13 ++++ CMakeLists.txt | 3 + cmake/pythonFormat.cmake | 25 ++++++++ python/edm4hep/__init__.py | 33 +++++----- scripts/createEDM4hepFile.py | 44 ++++---------- scripts/updateReadmeLinks.py | 96 +++++++++++++++++++----------- test/conftest.py | 2 +- test/test_EDM4hepFile.py | 52 ++++------------ test/test_rdf.py | 44 ++++++++------ test/tools/test_all_collections.py | 29 +++++---- test/utils/test_kinematics.py | 91 ++++++++++++++-------------- 13 files changed, 240 insertions(+), 204 deletions(-) create mode 100644 .ruff.toml create mode 100644 cmake/pythonFormat.cmake diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index e653c506f..7a035eb5e 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -32,7 +32,7 @@ jobs: cd $STARTDIR python3 -m venv /root/pre-commit-venv source /root/pre-commit-venv/bin/activate - pip install pre-commit + pip install pre-commit ruff cmake -DENABLE_SIO=ON \ -DCMAKE_CXX_STANDARD=20 \ -DCMAKE_CXX_FLAGS=" -fdiagnostics-color=always -Werror "\ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cbdbf8526..b54a59c8c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,3 +18,13 @@ repos: entry: python ./scripts/updateReadmeLinks.py --check files: (^README.md$) language: system + - id: ruff + name: ruff + entry: ruff check --force-exclude + types: [python] + language: system + - id: ruff-format + name: ruff-format + entry: ruff format --force-exclude + types: [python] + language: system diff --git a/.ruff.toml b/.ruff.toml new file mode 100644 index 000000000..98ef5e988 --- /dev/null +++ b/.ruff.toml @@ -0,0 +1,13 @@ +target-version = "py311" + +line-length = 99 + +[format] +# Make things format the same way as black +quote-style = "double" +indent-style = "space" +skip-magic-trailing-comma = false +line-ending = "auto" + +[lint] +select = ["F", "E", "W", "PLE", "PLW", "PLC"] diff --git a/CMakeLists.txt b/CMakeLists.txt index 19a52a3d7..bff3b3311 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -73,3 +73,6 @@ add_subdirectory(test) #--- create uninstall target --------------------------------------------------- include(cmake/EDM4HEPUninstall.cmake) + +#--- code format targets ------------------------------------------------------- +include(cmake/pythonFormat.cmake) diff --git a/cmake/pythonFormat.cmake b/cmake/pythonFormat.cmake new file mode 100644 index 000000000..d866d24fb --- /dev/null +++ b/cmake/pythonFormat.cmake @@ -0,0 +1,25 @@ +# Additional target to run python linters and formatters on python files +# +# Requires ruff to be available in the environment + +# Get all our Python files +file(GLOB_RECURSE ALL_PYTHON_FILES ${PROJECT_SOURCE_DIR}/python/*.py + ${PROJECT_SOURCE_DIR}/scripts/*.py ${PROJECT_SOURCE_DIR}/test/*.py) + +find_program(RUFF_EXECUTABLE ruff) +if(RUFF_EXECUTABLE) + add_custom_target( + ruff + COMMAND ${RUFF_EXECUTABLE} + check --force-exclude ${ALL_PYTHON_FILES} + ) + set_target_properties(ruff PROPERTIES EXCLUDE_FROM_ALL TRUE) + add_custom_target( + ruff-format + COMMAND ${RUFF_EXECUTABLE} + format --force-exclude ${ALL_PYTHON_FILES} + ) + set_target_properties(ruff-format PROPERTIES EXCLUDE_FROM_ALL TRUE) +else() + message(STATUS "Failed to find ruff executable - no target to run ruff can be set") +endif() diff --git a/python/edm4hep/__init__.py b/python/edm4hep/__init__.py index 205f4749c..8880cdf59 100644 --- a/python/edm4hep/__init__.py +++ b/python/edm4hep/__init__.py @@ -1,34 +1,37 @@ """Python module for the EDM4HEP datamodel and utilities.""" + import sys from .__version__ import __version__ import ROOT -res = ROOT.gSystem.Load('libedm4hepDict') + +res = ROOT.gSystem.Load("libedm4hepDict") if res < 0: - raise RuntimeError('Failed to load libedm4hepDict') + raise RuntimeError("Failed to load libedm4hepDict") -res = ROOT.gSystem.Load('libedm4hepRDF') +res = ROOT.gSystem.Load("libedm4hepRDF") if res < 0: - raise RuntimeError('Failed to load libedm4hepRDF') + raise RuntimeError("Failed to load libedm4hepRDF") -res = ROOT.gInterpreter.LoadFile('edm4hep/utils/kinematics.h') +res = ROOT.gInterpreter.LoadFile("edm4hep/utils/kinematics.h") if res != 0: - raise RuntimeError('Failed to load kinematics.h') + raise RuntimeError("Failed to load kinematics.h") -res = ROOT.gInterpreter.LoadFile('edm4hep/utils/dataframe.h') +res = ROOT.gInterpreter.LoadFile("edm4hep/utils/dataframe.h") if res != 0: - raise RuntimeError('Failed to load dataframe.h') + raise RuntimeError("Failed to load dataframe.h") -res = ROOT.gInterpreter.LoadFile('edm4hep/Constants.h') +res = ROOT.gInterpreter.LoadFile("edm4hep/Constants.h") if res != 0: - raise RuntimeError('Failed to load Constants.h') -from ROOT import edm4hep + raise RuntimeError("Failed to load Constants.h") +from ROOT import edm4hep # noqa: E402 + +from podio.pythonizations import load_pythonizations # noqa: E402 -from podio.pythonizations import load_pythonizations -load_pythonizations('edm4hep') +load_pythonizations("edm4hep") # Make TAB completion work for utils -setattr(edm4hep, 'utils', edm4hep.utils) +setattr(edm4hep, "utils", edm4hep.utils) # set package attributes for edm4hep edm4hep.__version__ = __version__ @@ -38,4 +41,4 @@ edm4hep.__file__ = __file__ # Make `import edm4hep` work -sys.modules['edm4hep'] = edm4hep +sys.modules["edm4hep"] = edm4hep diff --git a/scripts/createEDM4hepFile.py b/scripts/createEDM4hepFile.py index b532e174c..e7c2af07c 100755 --- a/scripts/createEDM4hepFile.py +++ b/scripts/createEDM4hepFile.py @@ -45,15 +45,9 @@ def create_MCParticleCollection(): particle.setCharge(next(counter)) particle.setTime(next(counter)) particle.setMass(next(counter)) - particle.setVertex( - edm4hep.Vector3d(next(counter), next(counter), next(counter)) - ) - particle.setEndpoint( - edm4hep.Vector3d(next(counter), next(counter), next(counter)) - ) - particle.setMomentum( - edm4hep.Vector3d(next(counter), next(counter), next(counter)) - ) + particle.setVertex(edm4hep.Vector3d(next(counter), next(counter), next(counter))) + particle.setEndpoint(edm4hep.Vector3d(next(counter), next(counter), next(counter))) + particle.setMomentum(edm4hep.Vector3d(next(counter), next(counter), next(counter))) particle.setMomentumAtEndpoint( edm4hep.Vector3d(next(counter), next(counter), next(counter)) ) @@ -174,9 +168,7 @@ def create_ClusterCollection(vectorsize, calo_hit): cluster.setPositionError(create_CovMatrixNf(3)) cluster.setITheta(next(counter)) cluster.setPhi(next(counter)) - cluster.setDirectionError( - edm4hep.Vector3f(next(counter), next(counter), next(counter)) - ) + cluster.setDirectionError(edm4hep.Vector3f(next(counter), next(counter), next(counter))) for j in range(vectorsize): cluster.addToShapeParameters(next(counter)) cluster.addToSubdetectorEnergies(next(counter)) @@ -256,9 +248,7 @@ def create_TrackCollection(vectorsize, tracker_hit): state.Z0 = next(counter) state.tanLambda = next(counter) state.time = next(counter) - state.referencePoint = edm4hep.Vector3f( - next(counter), next(counter), next(counter) - ) + state.referencePoint = edm4hep.Vector3f(next(counter), next(counter), next(counter)) state.covMatrix = create_CovMatrixNf(6) track.addToTrackStates(state) @@ -292,9 +282,7 @@ def create_ReconstructedParticleCollection(vertex, cluster, track): rparticle.setPDG(next(counter)) rparticle.setEnergy(next(counter)) rparticle.setMomentum(edm4hep.Vector3f(next(counter), next(counter), next(counter))) - rparticle.setReferencePoint( - edm4hep.Vector3f(next(counter), next(counter), next(counter)) - ) + rparticle.setReferencePoint(edm4hep.Vector3f(next(counter), next(counter), next(counter))) rparticle.setCharge(next(counter)) rparticle.setMass(next(counter)) rparticle.setGoodnessOfPID(next(counter)) @@ -402,9 +390,7 @@ def create_frame(): particles = frame.put(create_MCParticleCollection(), "MCParticleCollection") particle = particles[0] - hits = frame.put( - create_SimTrackerHitCollection(particle), "SimTrackerHitCollection" - ) + hits = frame.put(create_SimTrackerHitCollection(particle), "SimTrackerHitCollection") simtracker_hit = hits[0] calo_contribs = frame.put( @@ -422,9 +408,7 @@ def create_frame(): hits = frame.put(create_CalorimeterHitCollection(), "CalorimeterHitCollection") calo_hit = hits[0] - clusters = frame.put( - create_ClusterCollection(VECTORSIZE, calo_hit), "ClusterCollection" - ) + clusters = frame.put(create_ClusterCollection(VECTORSIZE, calo_hit), "ClusterCollection") cluster = clusters[0] hits = frame.put(create_TrackerHit3DCollection(), "TrackerHit3DCollection") @@ -433,9 +417,7 @@ def create_frame(): frame.put(create_RawTimeSeriesCollection(VECTORSIZE), "RawTimeSeriesCollection") - tracks = frame.put( - create_TrackCollection(VECTORSIZE, tracker_hit), "TrackCollection" - ) + tracks = frame.put(create_TrackCollection(VECTORSIZE, tracker_hit), "TrackCollection") track = tracks[0] pids, pid = create_ParticleIDCollection(VECTORSIZE) @@ -456,9 +438,7 @@ def create_frame(): put_link_collection(frame, "RecoMCParticleLink", reco_particle, particle) put_link_collection(frame, "CaloHitSimCaloHitLink", calo_hit, simcalo_hit) - put_link_collection( - frame, "TrackerHitSimTrackerHitLink", tracker_hit, simtracker_hit - ) + put_link_collection(frame, "TrackerHitSimTrackerHitLink", tracker_hit, simtracker_hit) put_link_collection(frame, "CaloHitMCParticleLink", calo_hit, particle) put_link_collection(frame, "ClusterMCParticleLink", cluster, particle) put_link_collection(frame, "TrackMCParticleLink", track, particle) @@ -481,9 +461,7 @@ def create_frame(): parser.add_argument( "--rntuple", action="store_true", help="Use a ROOT ntuple instead of EDM4hep" ) - parser.add_argument( - "--output-file", type=str, help="Output file name", default="edm4hep.root" - ) + parser.add_argument("--output-file", type=str, help="Output file name", default="edm4hep.root") args = parser.parse_args() output_file = args.output_file diff --git a/scripts/updateReadmeLinks.py b/scripts/updateReadmeLinks.py index 3901a7938..4aed7c963 100755 --- a/scripts/updateReadmeLinks.py +++ b/scripts/updateReadmeLinks.py @@ -5,46 +5,70 @@ parser = argparse.ArgumentParser() parser.add_argument("inputfile", help="The README file to run on", default="README.md") -parser.add_argument("--check", action='store_true', help='Use if you want to just check the links without actually changing the file') +parser.add_argument( + "--check", + action="store_true", + help="Use if you want to just check the links without actually changing the file", +) args = parser.parse_args() + def check_readme_links(): - edm4hep_yaml_content = "" - with open("edm4hep.yaml") as edm4hep_yaml: - edm4hep_yaml_content = edm4hep_yaml.readlines() - - readme_content = "" - with open("README.md", "r") as readme: - readme_content = readme.readlines() - - new_readme_content = "" - - links_are_ok = True - for readme_line in readme_content: - edm4hep_objects = re.findall("\[(.*?)\]\(https:\/\/github.com\/key4hep\/EDM4hep\/blob\/main\/edm4hep\.yaml#L(\d+?)\)", readme_line) - if edm4hep_objects: - for edm4hep_object, original_line_number in edm4hep_objects: - regex = f".*edm4hep::{edm4hep_object} *: *" - # find in which line it appears in edm4hep.yaml - edm4hep_yaml_line_numbers_with_match = [str(i + 1) for i, line in enumerate(edm4hep_yaml_content) if re.match(regex, line)] - if len(edm4hep_yaml_line_numbers_with_match) != 1: - print(f"Error: failed to replace line number for {edm4hep_object}, either no or several matches were found in edm4hep.yaml with the regex '{regex}'") - sys.exit(1) - if edm4hep_yaml_line_numbers_with_match[0] != original_line_number: - links_are_ok = False - readme_line = readme_line.replace(f"[{edm4hep_object}](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L{original_line_number})", f"[{edm4hep_object}](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L{edm4hep_yaml_line_numbers_with_match[0]})") - print(f"{edm4hep_object} is wrongly linked (line {original_line_number} --> {edm4hep_yaml_line_numbers_with_match[0]})") - new_readme_content += readme_line - return links_are_ok, new_readme_content + edm4hep_yaml_content = "" + with open("edm4hep.yaml") as edm4hep_yaml: + edm4hep_yaml_content = edm4hep_yaml.readlines() + + readme_content = "" + with open("README.md", "r") as readme: + readme_content = readme.readlines() + + new_readme_content = "" + + links_are_ok = True + for readme_content_line in readme_content: + readme_line = readme_content_line + edm4hep_objects = re.findall( + r"\[(.*?)\]\(https:\/\/github.com\/key4hep\/EDM4hep\/blob\/main\/edm4hep\.yaml#L(\d+?)\)", + readme_line, + ) + if edm4hep_objects: + for edm4hep_object, original_line_number in edm4hep_objects: + regex = f".*edm4hep::{edm4hep_object} *: *" + # find in which line it appears in edm4hep.yaml + edm4hep_yaml_line_numbers_with_match = [ + str(i + 1) + for i, line in enumerate(edm4hep_yaml_content) + if re.match(regex, line) + ] + if len(edm4hep_yaml_line_numbers_with_match) != 1: + print( + f"Error: failed to replace line number for {edm4hep_object}, either no " + f"or several matches were found in edm4hep.yaml with the regex '{regex}'" + ) + sys.exit(1) + if edm4hep_yaml_line_numbers_with_match[0] != original_line_number: + links_are_ok = False + readme_line = readme_line.replace( + f"[{edm4hep_object}](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L{original_line_number})", + f"[{edm4hep_object}](https://github.com/key4hep/EDM4hep/blob/main/edm4hep.yaml#L{edm4hep_yaml_line_numbers_with_match[0]})", + ) + print( + f"{edm4hep_object} is wrongly linked (line {original_line_number} --> {edm4hep_yaml_line_numbers_with_match[0]})" # noqa: E501 + ) + new_readme_content += readme_line + return links_are_ok, new_readme_content + links_are_ok, new_readme_content = check_readme_links() if links_are_ok: - print("README.md links are fine, nothing to change") + print("README.md links are fine, nothing to change") else: - if not args.check: - with open("README.md", "w") as readme: - readme.write(new_readme_content) - print("README.md links updated.") - else: - print("README.md links should be updated. (Run ./scripts/updateReadmeLinks.py to fix them)") - sys.exit(1) + if not args.check: + with open("README.md", "w") as readme: + readme.write(new_readme_content) + print("README.md links updated.") + else: + print( + "README.md links should be updated. (Run ./scripts/updateReadmeLinks.py to fix them)" + ) + sys.exit(1) diff --git a/test/conftest.py b/test/conftest.py index 228cc9be1..d56bd388a 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -14,5 +14,5 @@ def pytest_addoption(parser): def pytest_configure(config): """This is a slighty hacky solution to make the pytest configuration available in test modules outside of fixtures""" - global options + global options # noqa: PLW0603 options = config.option diff --git a/test/test_EDM4hepFile.py b/test/test_EDM4hepFile.py index 85b76f43e..a4f31c44e 100644 --- a/test/test_EDM4hepFile.py +++ b/test/test_EDM4hepFile.py @@ -103,9 +103,7 @@ def check_cov_matrix(cov_matrix, n_dim): assert cov_matrix[i] == next(counter) -def test_basic_file_contents( - reader, events, expected_edm4hep_version, expected_podio_version -): +def test_basic_file_contents(reader, events, expected_edm4hep_version, expected_podio_version): """Make sure the basic file contents are OK""" assert len(events) == FRAMES assert reader.current_file_version("edm4hep") == expected_edm4hep_version @@ -152,9 +150,7 @@ def test_MCParticleCollection(event): assert particle.getMomentumAtEndpoint() == edm4hep.Vector3d( next(counter), next(counter), next(counter) ) - assert particle.getSpin() == edm4hep.Vector3f( - next(counter), next(counter), next(counter) - ) + assert particle.getSpin() == edm4hep.Vector3f(next(counter), next(counter), next(counter)) assert particle.getColorFlow() == edm4hep.Vector2i(next(counter), next(counter)) assert particles[0].getDaughters()[0] == particles[1] @@ -172,12 +168,8 @@ def test_SimTrackerHitCollection(event, particle): assert hit.getTime() == next(counter) assert hit.getPathLength() == next(counter) assert hit.getQuality() == next(counter) - assert hit.getPosition() == edm4hep.Vector3d( - next(counter), next(counter), next(counter) - ) - assert hit.getMomentum() == edm4hep.Vector3f( - next(counter), next(counter), next(counter) - ) + assert hit.getPosition() == edm4hep.Vector3d(next(counter), next(counter), next(counter)) + assert hit.getMomentum() == edm4hep.Vector3f(next(counter), next(counter), next(counter)) assert hit.getParticle() == particle @@ -191,9 +183,7 @@ def test_CaloHitContributionCollection(event, particle): assert hit.getPDG() == next(counter) assert hit.getEnergy() == next(counter) assert hit.getTime() == next(counter) - assert hit.getStepPosition() == edm4hep.Vector3f( - next(counter), next(counter), next(counter) - ) + assert hit.getStepPosition() == edm4hep.Vector3f(next(counter), next(counter), next(counter)) assert hit.getParticle() == particle @@ -206,9 +196,7 @@ def test_SimCalorimeterHitCollection(event): hit = hits[0] assert hit.getCellID() == next(counter) assert hit.getEnergy() == next(counter) - assert hit.getPosition() == edm4hep.Vector3f( - next(counter), next(counter), next(counter) - ) + assert hit.getPosition() == edm4hep.Vector3f(next(counter), next(counter), next(counter)) calo_contrib = event.get("CaloHitContributionCollection")[0] assert len(hit.getContributions()) == 1 @@ -236,9 +224,7 @@ def test_CalorimeterHitCollection(event): assert hit.getEnergy() == next(counter) assert hit.getEnergyError() == next(counter) assert hit.getTime() == next(counter) - assert hit.getPosition() == edm4hep.Vector3f( - next(counter), next(counter), next(counter) - ) + assert hit.getPosition() == edm4hep.Vector3f(next(counter), next(counter), next(counter)) assert hit.getType() == next(counter) @@ -268,9 +254,7 @@ def test_ClusterCollection(event): assert cluster.getType() == next(counter) assert cluster.getEnergy() == next(counter) assert cluster.getEnergyError() == next(counter) - assert cluster.getPosition() == edm4hep.Vector3f( - next(counter), next(counter), next(counter) - ) + assert cluster.getPosition() == edm4hep.Vector3f(next(counter), next(counter), next(counter)) check_cov_matrix(cluster.getPositionError(), 3) assert cluster.getITheta() == next(counter) assert cluster.getPhi() == next(counter) @@ -304,9 +288,7 @@ def test_TrackerHit3DCollection(event): assert hit.getTime() == next(counter) assert hit.getEDep() == next(counter) assert hit.getEDepError() == next(counter) - assert hit.getPosition() == edm4hep.Vector3d( - next(counter), next(counter), next(counter) - ) + assert hit.getPosition() == edm4hep.Vector3d(next(counter), next(counter), next(counter)) check_cov_matrix(hit.getCovMatrix(), 3) @@ -326,9 +308,7 @@ def test_TrackerHitPlaneCollection(event): assert hit.getV() == edm4hep.Vector2f(next(counter), next(counter)) assert hit.getDu() == next(counter) assert hit.getDv() == next(counter) - assert hit.getPosition() == edm4hep.Vector3d( - next(counter), next(counter), next(counter) - ) + assert hit.getPosition() == edm4hep.Vector3d(next(counter), next(counter), next(counter)) check_cov_matrix(hit.getCovMatrix(), 3) @@ -401,9 +381,7 @@ def test_VertexCollection(event, reco_particle): assert v.getType() == next(counter) assert v.getChi2() == next(counter) assert v.getNdf() == next(counter) - assert v.getPosition() == edm4hep.Vector3f( - next(counter), next(counter), next(counter) - ) + assert v.getPosition() == edm4hep.Vector3f(next(counter), next(counter), next(counter)) check_cov_matrix(v.getCovMatrix(), 3) assert v.getAlgorithmType() == next(counter) assert len(v.getParameters()) == VECTORSIZE @@ -422,9 +400,7 @@ def test_ReconstructedParticleCollection(event, track): rparticle = rparticles[0] assert rparticle.getPDG() == next(counter) assert rparticle.getEnergy() == next(counter) - assert rparticle.getMomentum() == edm4hep.Vector3f( - next(counter), next(counter), next(counter) - ) + assert rparticle.getMomentum() == edm4hep.Vector3f(next(counter), next(counter), next(counter)) assert rparticle.getReferencePoint() == edm4hep.Vector3f( next(counter), next(counter), next(counter) ) @@ -538,9 +514,7 @@ def test_LinkCollections(event, particle, reco_particle, track): check_LinkCollection(event, "RecoMCParticleLink", reco_particle, particle) check_LinkCollection(event, "CaloHitSimCaloHitLink", calo_hit, simcalo_hit) - check_LinkCollection( - event, "TrackerHitSimTrackerHitLink", tracker_hit, simtracker_hit - ) + check_LinkCollection(event, "TrackerHitSimTrackerHitLink", tracker_hit, simtracker_hit) check_LinkCollection(event, "CaloHitMCParticleLink", calo_hit, particle) check_LinkCollection(event, "ClusterMCParticleLink", cluster, particle) check_LinkCollection(event, "TrackMCParticleLink", track, particle) diff --git a/test/test_rdf.py b/test/test_rdf.py index 7cdd59d9f..e28f5be59 100644 --- a/test/test_rdf.py +++ b/test/test_rdf.py @@ -5,25 +5,31 @@ ROOT.EnableImplicitMT() -print('Create RDataFrame ...') -df = ROOT.RDataFrame('events', 'edm4hep_events.root') +print("Create RDataFrame ...") +df = ROOT.RDataFrame("events", "edm4hep_events.root") -print('Apply selectors and define new branches ...') -df2 = (df.Define('MCParticles_pt', 'edm4hep::utils::pt(MCParticles)') - .Define('MCParticles_eta', 'edm4hep::utils::eta(MCParticles)') - .Define('MCParticles_cosTheta', 'edm4hep::utils::cos_theta(MCParticles)') - .Define('SimTrackerHits_r', 'edm4hep::utils::r(SimTrackerHits)') - .Define('SimTrackerHits_pt', 'edm4hep::utils::pt(SimTrackerHits)') - .Define('TrackerHits_r', 'edm4hep::utils::r(TrackerHitPlanes)') - ) +print("Apply selectors and define new branches ...") +df2 = ( + df.Define("MCParticles_pt", "edm4hep::utils::pt(MCParticles)") + .Define("MCParticles_eta", "edm4hep::utils::eta(MCParticles)") + .Define("MCParticles_cosTheta", "edm4hep::utils::cos_theta(MCParticles)") + .Define("SimTrackerHits_r", "edm4hep::utils::r(SimTrackerHits)") + .Define("SimTrackerHits_pt", "edm4hep::utils::pt(SimTrackerHits)") + .Define("TrackerHits_r", "edm4hep::utils::r(TrackerHitPlanes)") +) -filename = 'edm4hep_events_py_rdf.root' -print(f'Writing snapshot to disk ... \t{filename}') +filename = "edm4hep_events_py_rdf.root" +print(f"Writing snapshot to disk ... \t{filename}") -df2.Snapshot('events', filename, - ['MCParticles_pt', - 'MCParticles_eta', - 'MCParticles_cosTheta', - 'SimTrackerHits_r', - 'SimTrackerHits_pt', - 'TrackerHits_r']) +df2.Snapshot( + "events", + filename, + [ + "MCParticles_pt", + "MCParticles_eta", + "MCParticles_cosTheta", + "SimTrackerHits_r", + "SimTrackerHits_pt", + "TrackerHits_r", + ], +) diff --git a/test/tools/test_all_collections.py b/test/tools/test_all_collections.py index 1fbdd3e6d..adb1ea8bd 100644 --- a/test/tools/test_all_collections.py +++ b/test/tools/test_all_collections.py @@ -1,6 +1,6 @@ -''' +""" Tests if all datatypes are used in the cxx file. -''' +""" import sys import re @@ -9,40 +9,39 @@ def test(yamlfile_path, cxxfile_path): - ''' + """ Test itself. Takes two parameters, Podio YAML file location and cxx file to be checked. - ''' + """ - with open(yamlfile_path, mode='r', encoding="utf-8") as yamlfile: + with open(yamlfile_path, mode="r", encoding="utf-8") as yamlfile: datamodel = yaml.safe_load(yamlfile) # List stores lines of cxx code on which `insertToJson` is used datatypes_found = [] - with open(cxxfile_path, mode='r', encoding="utf-8") as cxxfile: + with open(cxxfile_path, mode="r", encoding="utf-8") as cxxfile: for cxxline in cxxfile.readlines(): - result = re.search('insertIntoJson', - cxxline) + result = re.search("insertIntoJson", cxxline) if result: - datatypes_found += ['edm4hep::' + result.group(1)] + datatypes_found += ["edm4hep::" + result.group(1)] datatypes_found = set(datatypes_found) - datatypes = set(datamodel['datatypes']) + datatypes = set(datamodel["datatypes"]) if not datatypes.issubset(datatypes_found): missing_datatypes = datatypes - datatypes_found - print('ERROR: One or more datatypes are not being converted:') + print("ERROR: One or more datatypes are not being converted:") for datatype in missing_datatypes: - print(' ' + datatype) + print(" " + datatype) sys.exit(2) if __name__ == "__main__": - parser = argparse.ArgumentParser(description='Test all collections') - parser.add_argument('yamlfile') - parser.add_argument('cxxfile') + parser = argparse.ArgumentParser(description="Test all collections") + parser.add_argument("yamlfile") + parser.add_argument("cxxfile") args = parser.parse_args() test(args.yamlfile, args.cxxfile) diff --git a/test/utils/test_kinematics.py b/test/utils/test_kinematics.py index f8e025cda..f460ecde0 100644 --- a/test/utils/test_kinematics.py +++ b/test/utils/test_kinematics.py @@ -1,7 +1,6 @@ #!/usr/bin/env python3 import unittest -import cppyy import edm4hep @@ -16,47 +15,49 @@ class TestKinematics(unittest.TestCase): - """Tests to ensure that the kinematics header only library also works in - python""" - - def test_p4(self): - """Test the p4 functionality since that has shown to trip-up cppyy""" - # podio doesn't expose the Mutable constructors so we have to use the full - # glory of the immutable types here - p = edm4hep.MCParticle(11, # PDG - -1, # generatorStatus - -1, # simulatorStatus - 1.0, # charge - 0.0, # time - 125.0, # charge - edm4hep.Vector3d(0, 0, 0), # vertex - edm4hep.Vector3d(0, 0, 0), # endpoint - edm4hep.Vector3d(1.0, 2.0, 3.0), # momentum - edm4hep.Vector3d(0, 0, 0), # momentumAtEndpoint - edm4hep.Vector3f(0, 0, 0), # spin - edm4hep.Vector2i(0, 0) # colorFlow - ) - - self.assertEqual(p4(p), LVM(1.0, 2.0, 3.0, 125.0)) - self.assertEqual(p4(p, SetEnergy(42.0)), LVE(1.0, 2.0, 3.0, 42.0)) - self.assertEqual(p4(p, SetMass(250.0)), LVM(1.0, 2.0, 3.0, 250.0)) - - p = edm4hep.ReconstructedParticle(11, # type - 250.0, # energy - edm4hep.Vector3f(1.0, 2.0, 3.0), # momentum - edm4hep.Vector3f(0, 0, 0), # referencePoint - 1.0, # charge - 125.0, # mass - 0.0, # goodnessOfPID - edm4hep.CovMatrix4f() # covMatrix - ) - - self.assertEqual(p4(p), LVM(1.0, 2.0, 3.0, 125.0)) - self.assertEqual(p4(p, UseEnergy), LVE(1.0, 2.0, 3.0, 250.0)) - self.assertEqual(p4(p, UseMass), LVM(1.0, 2.0, 3.0, 125.0)) - self.assertEqual(p4(p, SetMass()), LVM(1.0, 2.0, 3.0, 0.0)) - self.assertEqual(p4(p, SetEnergy(42.0)), LVE(1.0, 2.0, 3.0, 42.0)) - - -if __name__ == '__main__': - unittest.main() + """Tests to ensure that the kinematics header only library also works in + python""" + + def test_p4(self): + """Test the p4 functionality since that has shown to trip-up cppyy""" + # podio doesn't expose the Mutable constructors so we have to use the full + # glory of the immutable types here + p = edm4hep.MCParticle( + 11, # PDG + -1, # generatorStatus + -1, # simulatorStatus + 1.0, # charge + 0.0, # time + 125.0, # charge + edm4hep.Vector3d(0, 0, 0), # vertex + edm4hep.Vector3d(0, 0, 0), # endpoint + edm4hep.Vector3d(1.0, 2.0, 3.0), # momentum + edm4hep.Vector3d(0, 0, 0), # momentumAtEndpoint + edm4hep.Vector3f(0, 0, 0), # spin + edm4hep.Vector2i(0, 0), # colorFlow + ) + + self.assertEqual(p4(p), LVM(1.0, 2.0, 3.0, 125.0)) + self.assertEqual(p4(p, SetEnergy(42.0)), LVE(1.0, 2.0, 3.0, 42.0)) + self.assertEqual(p4(p, SetMass(250.0)), LVM(1.0, 2.0, 3.0, 250.0)) + + p = edm4hep.ReconstructedParticle( + 11, # type + 250.0, # energy + edm4hep.Vector3f(1.0, 2.0, 3.0), # momentum + edm4hep.Vector3f(0, 0, 0), # referencePoint + 1.0, # charge + 125.0, # mass + 0.0, # goodnessOfPID + edm4hep.CovMatrix4f(), # covMatrix + ) + + self.assertEqual(p4(p), LVM(1.0, 2.0, 3.0, 125.0)) + self.assertEqual(p4(p, UseEnergy), LVE(1.0, 2.0, 3.0, 250.0)) + self.assertEqual(p4(p, UseMass), LVM(1.0, 2.0, 3.0, 125.0)) + self.assertEqual(p4(p, SetMass()), LVM(1.0, 2.0, 3.0, 0.0)) + self.assertEqual(p4(p, SetEnergy(42.0)), LVE(1.0, 2.0, 3.0, 42.0)) + + +if __name__ == "__main__": + unittest.main()