Replace os with pathlib (#30)

Author: henry2004y

src/flekspy/__init__.py

@@ -2,38 +2,39 @@
 flekspy Public API.
 """
 
-import glob
-import os
+from pathlib import Path
 import errno
 from flekspy.idl import IDLData
 from flekspy.yt import FLEKSData, extract_phase
 from flekspy.tp import FLEKSTP
 
 
-def load(filename: str, iDomain=0, iSpecies=0, readFieldData=False):
+def load(filename: str, iDomain=0, iSpecies=0, readFieldData: bool = False):
     """Load FLEKS data.
 
     Args:
-        filename (str): Input file name.
+        filename (str): Input file name pattern.
         iDomain (int, optional): Test particle domain index. Defaults to 0.
         iSpecies (int, optional): Test particle species index. Defaults to 0.
         readFieldData (bool, optional): Whether or not to read field data for test particles. Defaults to False.
 
     Returns:
         FLEKS data: IDLData, FLEKSData, or FLEKSTP
     """
-    files = glob.glob(filename)
+    files = list(Path().glob(filename))
     if len(files) == 0:
-        raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), files)
-    filename = files[0]
+        message = f"No files found matching pattern: '{filename}'"
+        raise FileNotFoundError(errno.ENOENT, message, filename)
+    filename = str(files[0].resolve())
 
-    basename = os.path.basename(os.path.normpath(filename))
+    filepath = Path(filename)
+    basename = filepath.name
 
     if basename == "test_particles":
         return FLEKSTP(filename, iDomain=iDomain, iSpecies=iSpecies)
-    elif basename.find(".") != -1 and basename.split(".")[-1] in ["out", "outs"]:
+    elif filepath.suffix in [".out", ".outs"]:
         return IDLData(filename)
-    elif basename[-6:] == "_amrex":
+    elif basename.endswith("_amrex"):
         return FLEKSData(filename, readFieldData)
     else:
         raise Exception("Error: unknown file format!")

src/flekspy/tp/test_particles.py

@@ -1,7 +1,7 @@
 from typing import List, Tuple, Dict, Union, Callable
 
 import matplotlib.pyplot as plt
-import os
+from pathlib import Path
 import numpy as np
 import glob
 import struct
@@ -51,8 +51,8 @@ def __init__(
         if type(dirs) == str:
             dirs = [dirs]
 
-        header = dirs[0] + "/Header"
-        if os.path.exists(header):
+        header = Path(dirs[0] + "/Header")
+        if header.exists():
             with open(header, "r") as f:
                 self.nReal = int(f.readline())
         else:
@@ -88,8 +88,8 @@ def __init__(
         self.pfiles = self.pfiles[iListStart:iListEnd]
 
         self.plists: List[Dict[Tuple[int, int], int]] = []
-        for fileName in self.plistfiles:
-            self.plists.append(self.read_particle_list(fileName))
+        for filename in self.plistfiles:
+            self.plists.append(self.read_particle_list(filename))
 
         self.IDs = set()
         for plist in self.plists:
@@ -122,28 +122,28 @@ def get_index_to_time(self) -> List:
             print("Index to time mapping was not initialized")
         return self.indextotime
 
-    def read_particle_list(self, fileName: str) -> Dict[Tuple[int, int], int]:
+    def read_particle_list(self, filename: str) -> Dict[Tuple[int, int], int]:
         """
         Read and return a list of the particle IDs.
         """
         # 2 integers + 1 unsigned long long
         listUnitSize = 2 * 4 + 8
-        nByte = os.path.getsize(fileName)
+        nByte = Path(filename).stat().st_size
         nPart = int(nByte / listUnitSize)
         plist = {}
-        with open(fileName, "rb") as f:
+        with open(filename, "rb") as f:
             for _ in range(nPart):
                 binaryData = f.read(listUnitSize)
                 (cpu, id, loc) = struct.unpack("iiQ", binaryData)
                 plist.update({(cpu, id): loc})
         return plist
 
-    def _read_the_first_record(self, fileName: str) -> Union[List[float], None]:
+    def _read_the_first_record(self, filename: str) -> Union[List[float], None]:
         """
         Get the first record stored in one file.
         """
         dataList = list()
-        with open(fileName, "rb") as f:
+        with open(filename, "rb") as f:
             while True:
                 binaryData = f.read(4 * 4)
 
@@ -184,11 +184,11 @@ def read_particles_at_time(
                 break
             iFile += 1
 
-        fileName = self.pfiles[iFile]
+        filename = self.pfiles[iFile]
 
         dataList: list[float] = []
         idList: list[tuple] = []
-        with open(fileName, "rb") as f:
+        with open(filename, "rb") as f:
             while True:
                 binaryData = f.read(4 * 4)
                 if not binaryData:
@@ -217,14 +217,14 @@ def read_particles_at_time(
             raise Exception(f"There are no particles at time {time}.")
 
         if doSave:
-            fileName = f"particles_t{time}.csv"
+            filename = f"particles_t{time}.csv"
             header = "cpu,iid,time,x,y,z,ux,uy,uz"
             if self.nReal == 10:
                 header += ",bx,by,bz"
             elif self.nReal == 13:
                 header += ",bx,by,bz,ex,ey,ez"
 
-            with open(fileName, "w") as f:
+            with open(filename, "w") as f:
                 f.write(header + "\n")
                 for id_row, data_row in zip(idData, npData):
                     f.write(
@@ -236,7 +236,7 @@ def read_particles_at_time(
     def save_trajectory_to_csv(
         self,
         pID: Tuple[int, int],
-        fileName: str = None,
+        filename: str = None,
         shiftTime: bool = False,
         scaleTime: bool = False,
     ) -> None:
@@ -252,8 +252,8 @@ def save_trajectory_to_csv(
         >>> tp.save_trajectory_to_csv((3,15))
         """
         pData = self.read_particle_trajectory(pID)
-        if fileName == None:
-            fileName = "trajectory_" + str(pID[0]) + "_" + str(pID[1]) + ".csv"
+        if filename == None:
+            filename = "trajectory_" + str(pID[0]) + "_" + str(pID[1]) + ".csv"
         header = "time [s], X [R], Y [R], Z [R], U_x [km/s], U_y [km/s], U_z [km/s]"
         if self.nReal == 10:
             header += ", B_x [nT], B_y [nT], B_z [nT]"
@@ -265,7 +265,7 @@ def save_trajectory_to_csv(
             pData[:, 0] -= pData[0, 0]
             if scaleTime:
                 pData[:, 0] /= pData[-1, 0]
-        np.savetxt(fileName, pData, delimiter=",", header=header, comments="")
+        np.savetxt(filename, pData, delimiter=",", header=header, comments="")
 
     def read_particle_trajectory(self, pID: Tuple[int, int]):
         """
@@ -278,10 +278,10 @@ def read_particle_trajectory(self, pID: Tuple[int, int]):
         >>> trajectory = tp.read_particle_trajectory((66,888))
         """
         dataList = list()
-        for fileName, plist in zip(self.pfiles, self.plists):
+        for filename, plist in zip(self.pfiles, self.plists):
             if pID in plist:
                 ploc = plist[pID]
-                with open(fileName, "rb") as f:
+                with open(filename, "rb") as f:
                     f.seek(ploc)
                     binaryData = f.read(4 * 4)
                     (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)
@@ -298,10 +298,10 @@ def read_initial_location(self, pID):
         Return the initial location of a test particle.
         """
 
-        for fileName, plist in zip(self.pfiles, self.plists):
+        for filename, plist in zip(self.pfiles, self.plists):
             if pID in plist:
                 ploc = plist[pID]
-                with open(fileName, "rb") as f:
+                with open(filename, "rb") as f:
                     f.seek(ploc)
                     binaryData = f.read(4 * 4)
                     (cpu, idtmp, nRecord, weight) = struct.unpack("iiif", binaryData)

src/flekspy/util/utilities.py

@@ -135,26 +135,29 @@ def download_testfile(url: str, target_path="."):
       url (str): the URL of the tar.gz file.
       target_path (str): the directory to extract the files to. Defaults to the current directory.
     """
-    import os, requests
+    import requests
+    from pathlib import Path
+
+    target_dir = Path(target_path)
+    temp_file = Path("temp.tar.gz")
 
     try:
         response = requests.get(url, stream=True)
-        response.raise_for_status()  # Raise an exception for bad status codes
+        response.raise_for_status() # Raise an exception for bad status codes
 
-        with open("temp.tar.gz", "wb") as f:
+        with open(temp_file, "wb") as f:
             for chunk in response.iter_content(chunk_size=8192):
                 f.write(chunk)
 
-        if not os.path.exists(target_path):
-            os.makedirs(target_path)
+        target_dir.mkdir(parents=True, exist_ok=True)
 
-        with tarfile.open("temp.tar.gz", "r:gz") as tar:
-            tar.extractall(target_path)
+        with tarfile.open(temp_file, "r:gz") as tar:
+            tar.extractall(target_dir)
 
     except requests.exceptions.RequestException as e:
         print(f"Error downloading file: {e}")
     except tarfile.TarError as e:
         print(f"Error extracting tar file: {e}")
     finally:
-        if os.path.exists("temp.tar.gz"):
-            os.remove("temp.tar.gz")  # Clean up temporary file
+        if temp_file.exists():
+            temp_file.unlink(missing_ok=True) # Clean up temporary file

src/flekspy/yt/yt.py

@@ -1,5 +1,4 @@
-import os
-import glob
+from pathlib import Path
 import numpy as np
 
 import yt
@@ -159,7 +158,7 @@ def __init__(
     def _parse_parameter_file(self):
         super(FLEKSData, self)._parse_parameter_file()
 
-        fleks_header = os.path.join(self.output_dir, "FLEKSHeader")
+        fleks_header = Path(self.output_dir) / "FLEKSHeader"
         with open(fleks_header, "r") as f:
             plot_string = f.readline().lower()
             self.radius = float(f.readline())  # should be in unit [m]
@@ -178,8 +177,10 @@ def _parse_parameter_file(self):
         else:
             self.parameters["fleks_unit"] = "unknown"
 
-        particle_types = glob.glob(self.output_dir + "/*/Header")
-        particle_types = [cpt.split(os.sep)[-2] for cpt in particle_types]
+        output_dir_path = Path(self.output_dir)
+        header_paths_generator = output_dir_path.glob("*/Header")
+        particle_types = [p.parent.name for p in header_paths_generator]
+
         if len(particle_types) > 0 and not self.read_field_data:
             self.parameters["particles"] = 1
             self.particle_types = tuple(particle_types)

tests/test_fleks.py

@@ -26,12 +26,12 @@
 
 
 class TestIDL:
-    files = (
+    filenames = (
         "1d__raw_2_t25.60000_n00000258.out",
         "z=0_fluid_region0_0_t00001640_n00010142.out",
         "3d_raw.out",
     )
-    files = [os.path.join("tests/data/", file) for file in files]
+    files = [os.path.join("tests/data/", file) for file in filenames]
 
     def test_load(self):
         ds = fs.load(self.files[0])
@@ -91,9 +91,14 @@ def test_phase(self):
 
         ## Select and plot the particles inside a box defined by xleft and xright
         region = ds.box(xleft, xright)
-        x, y, w = ds.get_phase(x_field, y_field, z_field, region=region,
-            domain_size=(xleft[0], xright[0], xleft[1], xright[1]))
-        assert x.shape == (128,) and w.max() == 2.8024863240162035e+23
+        x, y, w = ds.get_phase(
+            x_field,
+            y_field,
+            z_field,
+            region=region,
+            domain_size=(xleft[0], xright[0], xleft[1], xright[1]),
+        )
+        assert x.shape == (128,) and w.max() == 2.8024863240162035e23
         pp = ds.plot_phase(
             x_field,
             y_field,
@@ -183,12 +188,11 @@ def load(files):
 
 
 def test_load(benchmark):
-    path = "tests/data/"
-    files = (
+    filenames = (
         "1d__raw_2_t25.60000_n00000258.out",
         "z=0_fluid_region0_0_t00001640_n00010142.out",
     )
-    files = [os.path.join(path, file) for file in files]
+    files = [os.path.join("tests/data/", file) for file in filenames]
 
     result = benchmark(load, files)