Merge pull request #704 from janosh/fix-zpath-pathlib

Fix `zpath` when passing `pathlib.Path`

src/monty/os/path.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import os
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 from monty.fnmatch import WildCard
@@ -14,7 +15,7 @@
  from typing import Callable, Literal, Optional, Union
 
 
-def zpath(filename: str) -> str:
+def zpath(filename: str | Path) -> str:
  """
  Returns an existing (zipped or unzipped) file path given the unzipped
  version. If no path exists, returns the filename unmodified.
@@ -23,10 +24,10 @@ def zpath(filename: str) -> str:
  filename: filename without zip extension
 
  Returns:
- filename with a zip extension (unless an unzipped version
- exists). If filename is not found, the same filename is returned
- unchanged.
+ str: filename with a zip extension (unless an unzipped version exists).
+ If filename is not found, the same filename is returned unchanged.
  """
+ filename = str(filename) # ensure we work with strings
  exts = ("", ".gz", ".GZ", ".bz2", ".BZ2", ".z", ".Z")
  for ext in exts:
  filename = filename.removesuffix(ext)

tests/test_os.py

@@ -13,17 +13,51 @@
 
 
 class TestPath:
- def test_zpath(self, tmp_path: Path):
+ def test_zpath_str(self, tmp_path: Path):
  tmp_gz = tmp_path / "tmp.gz"
  tmp_gz.touch()
  ret_path = zpath(str(tmp_gz))
+ assert isinstance(ret_path, str)
  assert ret_path == str(tmp_gz)
 
  tmp_not_bz2 = tmp_path / "tmp_not_bz2"
  tmp_not_bz2.touch()
 
  ret_path = zpath(f"{tmp_not_bz2}.bz2")
  assert ret_path == str(tmp_not_bz2)
+ assert isinstance(ret_path, str)
+
+ def test_zpath_path(self, tmp_path: Path):
+ # Test with Path input
+ tmp_gz = tmp_path / "tmp.gz"
+ tmp_gz.touch()
+ ret_path = zpath(tmp_gz)
+ assert ret_path == str(tmp_gz)
+ assert isinstance(ret_path, str)
+
+ def test_zpath_multiple_extensions(self, tmp_path: Path):
+ exts = ["", ".gz", ".GZ", ".bz2", ".BZ2", ".z", ".Z"]
+ for ext in exts:
+ tmp_file = tmp_path / f"tmp{ext}"
+ # create files with all supported compression extensions
+ tmp_file.touch()
+
+ # zpath should return the file without compression extension
+ ret_path = zpath(tmp_path / "tmp")
+ assert ret_path == str(tmp_path / "tmp")
+ assert isinstance(ret_path, str)
+
+ (tmp_path / "tmp").unlink() # Remove the uncompressed file
+ ret_path = zpath(tmp_path / "tmp")
+ assert ret_path == str(tmp_path / "tmp.gz") # should find .gz first now
+
+ def test_zpath_nonexistent_file(self, tmp_path: Path):
+ # should return path as is for non-existent file
+ nonexistent = tmp_path / "nonexistent.txt"
+ ret_path = zpath(nonexistent)
+ assert ret_path == str(nonexistent)
+ ret_path = zpath(f"{nonexistent}.bz2")
+ assert ret_path == str(nonexistent)
 
  def test_find_exts(self):
  assert len(find_exts(MODULE_DIR, "py")) >= 18