Repository: scikit-hep/pylhe

Issue:
Include ability to read gzip'ed input LHE file
As a quick suggestion for an improvement, it might be very handy to include the ability of opening zipped lhe files without the need for unzipping beforehand.
This can be achieved checking the file type and in case using something like
https://stackoverflow.com/questions/10566558/python-read-lines-from-compressed-text-files

This feature can be then disabled if the right dependency is not there, but it looks to me a rather minimal dependency and a quite wide usage for this feature.

 Thanks for considering it.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/scikit-hep/pylhe/pull/74?src=pr&el=h1) Report
> Merging [#74](https://codecov.io/gh/scikit-hep/pylhe/pull/74?src=pr&el=desc) (a1637e1) into [master](https://codecov.io/gh/scikit-hep/pylhe/commit/fa6968b80fcb3a90e0117b0b2e2bb8ede004f619?el=desc) (fa6968b) will **decrease** coverage by `1.78%`.
> The diff coverage is `45.00%`.

[![Impacted file tree graph](https://codecov.io/gh/scikit-hep/pylhe/pull/74/graphs/tree.svg?width=650&height=150&src=pr&token=I5YnvSyPCo)](https://codecov.io/gh/scikit-hep/pylhe/pull/74?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master      #74      +/-   ##
==========================================
- Coverage   62.80%   61.01%   -1.79%     
==========================================
  Files           1        1              
  Lines         164      177      +13     
  Branches       32       33       +1     
==========================================
+ Hits          103      108       +5     
- Misses         53       60       +7     
- Partials        8        9       +1     
```

| Flag | Coverage Δ | |
|---|---|---|
| unittests | `61.01% <45.00%> (-1.79%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/scikit-hep/pylhe/pull/74?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/pylhe/\_\_init\_\_.py](https://codecov.io/gh/scikit-hep/pylhe/pull/74/diff?src=pr&el=tree#diff-c3JjL3B5bGhlL19faW5pdF9fLnB5) | `61.01% <45.00%> (-1.79%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/scikit-hep/pylhe/pull/74?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/scikit-hep/pylhe/pull/74?src=pr&el=footer). Last update [fa6968b...a1637e1](https://codecov.io/gh/scikit-hep/pylhe/pull/74?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

matthewfeickert: @lukasheinrich I think this is finally ready for review (after me spamming this PR for the last hour). :P

Patch:
diff --git a/src/pylhe/__init__.py b/src/pylhe/__init__.py
index 1cf7ec0..821c8de 100644
--- a/src/pylhe/__init__.py
+++ b/src/pylhe/__init__.py
@@ -1,8 +1,10 @@
+import gzip
 import os
+import subprocess
 import xml.etree.ElementTree as ET
+
 import networkx as nx
 import tex2pix
-import subprocess
 from particle.converters.bimap import DirectionalMaps
 
 
@@ -120,101 +122,131 @@ def loads():
     pass
 
 
-def readLHEInit(thefile):
+def _extract_fileobj(filepath):
+    """
+    Checks to see if a file is compressed, and if so, extract it with gzip
+    so that the uncompressed file can be returned.
+    It returns a file object containing XML data that will be ingested by
+    ``xml.etree.ElementTree.iterparse``.
+
+    Args:
+        filepath: A path-like object or str.
+
+    Returns:
+        _io.BufferedReader or gzip.GzipFile: A file object containing XML data.
+    """
+    with open(filepath, "rb") as gzip_file:
+        header = gzip_file.read(2)
+    gzip_magic_number = b"\x1f\x8b"
+
+    return (
+        gzip.GzipFile(filepath) if header == gzip_magic_number else open(filepath, "rb")
+    )
+
+
+def readLHEInit(filepath):
     """
-    Read the init blocks. Return dict. This encodes the weight group
+    Read and return the init blocks. This encodes the weight group
     and related things according to https://arxiv.org/abs/1405.1067
-    This function returns a dict.
+
+    Args:
+        filepath: A path-like object or str.
+
+    Returns:
+        dict: Dictionary containing the init blocks of the LHE file.
     """
     initDict = {}
-    for event, element in ET.iterparse(thefile, events=["end"]):
-        if element.tag == "init":
-            data = element.text.split("\n")[1:-1]
-            initDict["initInfo"] = LHEInit.fromstring(data[0])
-            initDict["procInfo"] = [LHEProcInfo.fromstring(d) for d in data[1:]]
-        if element.tag == "initrwgt":
-            initDict["weightgroup"] = {}
-            for child in element:
-                # Find all weightgroups
-                if child.tag == "weightgroup" and child.attrib != {}:
-                    try:
-                        wg_type = child.attrib["type"]
-                    except KeyError:
-                        print("weightgroup must have attribute 'type'")
-                        raise
-                    _temp = {"attrib": child.attrib}
-                    _temp["weights"] = {}
-                    # Iterate over all weights in this weightgroup
-                    for w in child:
-                        if not w.tag == "weight":
-                            continue
+    with _extract_fileobj(filepath) as fileobj:
+        for event, element in ET.iterparse(fileobj, events=["end"]):
+            if element.tag == "init":
+                data = element.text.split("\n")[1:-1]
+                initDict["initInfo"] = LHEInit.fromstring(data[0])
+                initDict["procInfo"] = [LHEProcInfo.fromstring(d) for d in data[1:]]
+            if element.tag == "initrwgt":
+                initDict["weightgroup"] = {}
+                for child in element:
+                    # Find all weightgroups
+                    if child.tag == "weightgroup" and child.attrib != {}:
                         try:
-                            wg_id = w.attrib["id"]
+                            wg_type = child.attrib["type"]
                         except KeyError:
-                            print("weight must have attribute 'id'")
+                            print("weightgroup must have attribute 'type'")
                             raise
-                        _w = {"attrib": w.attrib}
-                        _w["name"] = w.text.strip()
-                        _temp["weights"][wg_id] = _w
-
-                    initDict["weightgroup"][wg_type] = _temp
-        if element.tag == "event":
-            break
+                        _temp = {"attrib": child.attrib}
+                        _temp["weights"] = {}
+                        # Iterate over all weights in this weightgroup
+                        for w in child:
+                            if not w.tag == "weight":
+                                continue
+                            try:
+                                wg_id = w.attrib["id"]
+                            except KeyError:
+                                print("weight must have attribute 'id'")
+                                raise
+                            _w = {"attrib": w.attrib}
+                            _w["name"] = w.text.strip()
+                            _temp["weights"][wg_id] = _w
+
+                        initDict["weightgroup"][wg_type] = _temp
+            if element.tag == "event":
+                break
     return initDict
 
 
-def readLHE(thefile):
+def readLHE(filepath):
     try:
-        for event, element in ET.iterparse(thefile, events=["end"]):
-            if element.tag == "event":
-                data = element.text.split("\n")[1:-1]
-                eventdata, particles = data[0], data[1:]
-                eventinfo = LHEEventInfo.fromstring(eventdata)
-                particle_objs = []
-                for p in particles:
-                    particle_objs += [LHEParticle.fromstring(p)]
-                yield LHEEvent(eventinfo, particle_objs)
+        with _extract_fileobj(filepath) as fileobj:
+            for event, element in ET.iterparse(fileobj, events=["end"]):
+                if element.tag == "event":
+                    data = element.text.split("\n")[1:-1]
+                    eventdata, particles = data[0], data[1:]
+                    eventinfo = LHEEventInfo.fromstring(eventdata)
+                    particle_objs = []
+                    for p in particles:
+                        particle_objs += [LHEParticle.fromstring(p)]
+                    yield LHEEvent(eventinfo, particle_objs)
     except ET.ParseError as e:
         print("WARNING. Parse Error:", e)
         return
 
 
-def readLHEWithAttributes(thefile):
+def readLHEWithAttributes(filepath):
     """
     Iterate through file, similar to readLHE but also set
     weights and attributes.
     """
     try:
-        for event, element in ET.iterparse(thefile, events=["end"]):
-            if element.tag == "event":
-                eventdict = {}
-                data = element.text.split("\n")[1:-1]
-                eventdata, particles = data[0], data[1:]
-                eventdict["eventinfo"] = LHEEventInfo.fromstring(eventdata)
-                eventdict["particles"] = []
-                eventdict["weights"] = {}
-                eventdict["attrib"] = element.attrib
-                eventdict["optional"] = []
-                for p in particles:
-                    if not p.strip().startswith("#"):
-                        eventdict["particles"] += [LHEParticle.fromstring(p)]
-                    else:
-                        eventdict["optional"].append(p.strip())
-                for sub in element:
-                    if sub.tag == "rwgt":
-                        for r in sub:
-                            if r.tag == "wgt":
-                                eventdict["weights"][r.attrib["id"]] = float(
-                                    r.text.strip()
-                                )
-                # yield eventdict
-                yield LHEEvent(
-                    eventdict["eventinfo"],
-                    eventdict["particles"],
-                    eventdict["weights"],
-                    eventdict["attrib"],
-                    eventdict["optional"],
-                )
+        with _extract_fileobj(filepath) as fileobj:
+            for event, element in ET.iterparse(fileobj, events=["end"]):
+                if element.tag == "event":
+                    eventdict = {}
+                    data = element.text.split("\n")[1:-1]
+                    eventdata, particles = data[0], data[1:]
+                    eventdict["eventinfo"] = LHEEventInfo.fromstring(eventdata)
+                    eventdict["particles"] = []
+                    eventdict["weights"] = {}
+                    eventdict["attrib"] = element.attrib
+                    eventdict["optional"] = []
+                    for p in particles:
+                        if not p.strip().startswith("#"):
+                            eventdict["particles"] += [LHEParticle.fromstring(p)]
+                        else:
+                            eventdict["optional"].append(p.strip())
+                    for sub in element:
+                        if sub.tag == "rwgt":
+                            for r in sub:
+                                if r.tag == "wgt":
+                                    eventdict["weights"][r.attrib["id"]] = float(
+                                        r.text.strip()
+                                    )
+                    # yield eventdict
+                    yield LHEEvent(
+                        eventdict["eventinfo"],
+                        eventdict["particles"],
+                        eventdict["weights"],
+                        eventdict["attrib"],
+                        eventdict["optional"],
+                    )
     except ET.ParseError:
         print("WARNING. Parse Error.")
         return


Test patch:
diff --git a/tests/test_lhe_reader.py b/tests/test_lhe_reader.py
index cdcebf5..aff4674 100644
--- a/tests/test_lhe_reader.py
+++ b/tests/test_lhe_reader.py
@@ -1,10 +1,51 @@
+import gzip
+import os
+import shutil
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+
 import pytest
-import pylhe
 import skhep_testdata
 
+import pylhe
+
 TEST_FILE = skhep_testdata.data_path("pylhe-testfile-pr29.lhe")
 
 
+@pytest.fixture(scope="session")
+def testdata_gzip_file():
+    test_data = skhep_testdata.data_path("pylhe-testfile-pr29.lhe")
+    tmp_path = Path(NamedTemporaryFile().name)
+
+    # create what is basically pylhe-testfile-pr29.lhe.gz
+    with open(test_data, "rb") as readfile:
+        with gzip.open(tmp_path, "wb") as writefile:
+            shutil.copyfileobj(readfile, writefile)
+    yield tmp_path
+
+    # teardown
+    os.remove(tmp_path)
+
+
+def test_gzip_open(tmpdir, testdata_gzip_file):
+    assert pylhe._extract_fileobj(TEST_FILE)
+    assert pylhe._extract_fileobj(testdata_gzip_file)
+
+    # Needs path-like object, not a fileobj
+    with pytest.raises(TypeError):
+        with open(TEST_FILE, "rb") as fileobj:
+            pylhe._extract_fileobj(fileobj)
+
+    with open(TEST_FILE, "rb") as fileobj:
+        assert isinstance(pylhe._extract_fileobj(TEST_FILE), type(fileobj))
+        assert isinstance(pylhe._extract_fileobj(Path(TEST_FILE)), type(fileobj))
+    assert isinstance(pylhe._extract_fileobj(testdata_gzip_file), gzip.GzipFile)
+    assert isinstance(pylhe._extract_fileobj(Path(testdata_gzip_file)), gzip.GzipFile)
+
+    # Verify uncompressed and compressed both work
+    assert pylhe.readLHEInit(TEST_FILE) == pylhe.readLHEInit(testdata_gzip_file)
+
+
 def test_event_count():
     assert pylhe.readNumEvents(TEST_FILE) == 791
 
@@ -22,3 +63,8 @@ def test_lhe_init():
     assert init_info["PDFsetB"] == pytest.approx(6.0)
     assert init_info["weightingStrategy"] == pytest.approx(7.0)
     assert init_info["numProcesses"] == pytest.approx(8.0)
+
+
+def test_readLHE(testdata_gzip_file):
+    assert pylhe.readLHE(TEST_FILE)
+    assert pylhe.readLHE(testdata_gzip_file)