feat: add LHEFile.count_events (#334)

* feat: LHEFile.count_events

Closes: #333

* Also terminate event parsing on closing LesHoucheEvents tag

* fix: typo

* Clean comments

* rename

* test count returns negative number of events

* print -> warnings for parse errors

* expect a RuntimeWarning

* Test for Parse Error in fromfile

* Make StopIteration error more understandable

Author: APN-Pucky

README.md

@@ -61,8 +61,8 @@ from skhep_testdata import data_path
 import pylhe
 
 lhe_file = data_path("pylhe-testlhef3.lhe")
-events = pylhe.read_lhe_with_attributes(lhe_file)
-print(f"Number of events: {pylhe.read_num_events(lhe_file)}")
+events = pylhe.LHEFile.fromfile(lhe_file).events
+print(f"Number of events: {pylhe.LHEFile.count_events(lhe_file)}")
 
 # Get event 1
 event = next(itertools.islice(events, 1, 2))

benchmarks/test_count_events_bench.py

@@ -0,0 +1,40 @@
+"""
+Benchmark tests for pylhe event counting performance.
+"""
+
+import skhep_testdata
+
+import pylhe
+
+# Test data files from skhep_testdata - all LHE and LHE.gz files
+TEST_FILES_LHE_ALL = [
+    skhep_testdata.data_path("pylhe-testfile-pr29.lhe"),
+    skhep_testdata.data_path("pylhe-testlhef3.lhe"),
+    *[
+        skhep_testdata.data_path(f"pylhe-testfile-powheg-box-v2-{proc}.lhe")
+        for proc in ["Z", "W", "Zj", "trijet", "directphoton", "hvq"]
+    ],
+    skhep_testdata.data_path("pylhe-testfile-madgraph-2.0.0-wbj.lhe"),
+    skhep_testdata.data_path("pylhe-testfile-madgraph-2.2.1-Z-ckkwl.lhe.gz"),
+    skhep_testdata.data_path("pylhe-testfile-madgraph-2.2.1-Z-fxfx.lhe.gz"),
+    skhep_testdata.data_path("pylhe-testfile-madgraph-2.2.1-Z-mlm.lhe.gz"),
+    skhep_testdata.data_path("pylhe-testfile-madgraph5-3.5.8-pp_to_jj.lhe.gz"),
+    skhep_testdata.data_path("pylhe-testfile-pythia-6.413-ttbar.lhe"),
+    skhep_testdata.data_path("pylhe-testfile-pythia-8.3.14-weakbosons.lhe"),
+    skhep_testdata.data_path("pylhe-testfile-sherpa-3.0.1-eejjj.lhe"),
+    skhep_testdata.data_path("pylhe-testfile-whizard-3.1.4-eeWW.lhe"),
+]
+
+
+def test_count_events_benchmark(benchmark):
+    """Benchmark using the count_events function across all test files."""
+
+    def count_events_all_files(filepaths):
+        total_events = 0
+        for filepath in filepaths:
+            num_events = pylhe.LHEFile.count_events(filepath)
+            total_events += num_events
+        return total_events
+
+    result = benchmark(count_events_all_files, TEST_FILES_LHE_ALL)
+    print(f"Total events across all files: {result}")

src/pylhe/__init__.py

@@ -692,6 +692,8 @@ def _fromcontext(
                 # Clear memory
                 element.clear()
                 root.clear()
+            if element.tag == "LesHouchesEvents" and event == "end":
+                return
 
     @property
     def graph(self) -> graphviz.Digraph:
@@ -860,7 +862,7 @@ def _generator(lhef: LHEFile) -> Iterator[LHEEvent]:
                     )
 
             except ET.ParseError as excep:
-                print("WARNING. Parse Error:", excep)
+                warnings.warn(f"Parse Error: {excep}", RuntimeWarning, stacklevel=1)
                 return
 
         lhef = cls(
@@ -874,11 +876,45 @@ def _generator(lhef: LHEFile) -> Iterator[LHEEvent]:
             events=[],
         )
         events = _generator(lhef)
-        next(events)  # advance to read lheinit
+        try:
+            next(events)  # advance to read lheinit
+        except StopIteration:
+            # If generator stops without yielding, it means no init was read
+            err = "No or faulty <init> block found in the LHE file."
+            raise ValueError(err) from None
 
         lhef.events = events if generator else list(events)
         return lhef
 
+    @staticmethod
+    def count_events(filepath: PathLike) -> int:
+        """
+        Efficiently count the number of events in an LHE file without loading them into memory.
+
+        Args:
+            filepath: Path to the LHE file.
+
+        Returns:
+            Number of events in the file, or -1 if parsing fails.
+        """
+        try:
+            with _extract_fileobj(filepath) as fileobj:
+                context = ET.iterparse(fileobj, events=["start", "end"])
+                _, root = next(context)  # Get the root element
+                count = 0
+                for event, element in context:
+                    if event == "end" and element.tag == "event":
+                        count += 1
+                        # Clear the element to free memory
+                        element.clear()
+                        # Root tracks sub-elements -> clear all sub-elements
+                        root.clear()
+                    if event == "end" and element.tag == "LesHouchesEvents":
+                        return count
+        except ET.ParseError as excep:
+            warnings.warn(f"Parse Error: {excep}", RuntimeWarning, stacklevel=1)
+        return -1
+
 
 def read_lhe_file(filepath: PathLike, with_attributes: bool = True) -> LHEFile:
     """
@@ -998,23 +1034,17 @@ def read_lhe_with_attributes(filepath: PathLike) -> Iterable[LHEEvent]:
 def read_num_events(filepath: PathLike) -> int:
     """
     Moderately efficient way to get the number of events stored in a file.
+
+    .. deprecated:: 1.0.0
+        Use `LHEFile.count_events` instead.
     """
-    try:
-        with _extract_fileobj(filepath) as fileobj:
-            context = ET.iterparse(fileobj, events=["start", "end"])
-            _, root = next(context)  # Get the root element
-            count = 0
-            for event, element in context:
-                if event == "end" and element.tag == "event":
-                    count += 1
-                    # Clear the element to free memory
-                    element.clear()
-                    # Root tracks sub-elements -> clear all sub-elements
-                    root.clear()
-            return count
-    except ET.ParseError as excep:
-        print("WARNING. Parse Error:", excep)
-        return -1
+    warnings.warn(
+        "read_num_events is deprecated and will be removed in a future version. "
+        "Use `LHEFile.count_events` instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return LHEFile.count_events(filepath)
 
 
 def write_lhe_file_string(

tests/test_errors.py

@@ -233,3 +233,43 @@ def test_whitespace_only_wgt_block_error():
             list(pylhe.read_lhe_with_attributes(tmp_file_path))
     finally:
         os.unlink(tmp_file_path)
+
+
+def test_count_events_parse_error():
+    """Test that ParseError warning is issued and -1 returned when counting events in malformed LHE file."""
+    # Create a temporary file with invalid XML content
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".lhe", delete=True) as f:
+        # Write invalid XML that will cause a parse error
+        f.write('<LesHouchesEvents version="3.0">\n')
+        f.write("<init>\n")
+        f.write("invalid xml content without proper closing\n")
+        # Missing </init> and </LesHouchesEvents> tags
+
+        f.flush()
+
+        # Test that a RuntimeWarning is issued and -1 is returned
+        with pytest.warns(RuntimeWarning, match=r"Parse Error:"):
+            assert pylhe.LHEFile.count_events(f.name) == -1
+
+
+def test_fromfile_parse_error():
+    """Test that ParseError warning is issued when loading malformed LHE file with fromfile."""
+    # Create a temporary file with invalid XML content
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".lhe", delete=True) as f:
+        # Write invalid XML that will cause a parse error
+        f.write('<LesHouchesEvents version="3.0">\n')
+        f.write("<init>\n")
+        f.write("invalid xml content without proper closing\n")
+        # Missing </init> and </LesHouchesEvents> tags
+
+        f.flush()
+
+        # Test that a RuntimeWarning is issued when trying to load the malformed file
+        # and potentially a ValueError if the generator stops without yielding
+        with (
+            pytest.warns(RuntimeWarning, match=r"Parse Error:"),
+            pytest.raises(
+                ValueError, match=r"No or faulty <init> block found in the LHE file"
+            ),
+        ):
+            pylhe.LHEFile.fromfile(f.name)

tests/test_lhe_reader.py

@@ -44,6 +44,11 @@
     *TEST_FILES_LHE_SHERPA,
     *TEST_FILES_LHE_WHIZARD,
 ]
+TEST_FILES_LHE_ALL = [
+    TEST_FILE_LHE_v1,
+    TEST_FILE_LHE_v3,
+    *TEST_FILES_LHE_GENERATORS,
+]
 
 
 @pytest.fixture(scope="session")
@@ -82,6 +87,13 @@ def test_read_num_events(testdata_gzip_file):
     )
 
 
+@pytest.mark.parametrize("file", TEST_FILES_LHE_ALL)
+def test_count_events(file):
+    assert pylhe.LHEFile.count_events(file) == sum(
+        1 for _ in pylhe.LHEFile.fromfile(file).events
+    )
+
+
 def test_read_lhe_init_gzipped_file(testdata_gzip_file):
     assert pylhe.read_lhe_init(TEST_FILE_LHE_v1) == pylhe.read_lhe_init(
         testdata_gzip_file