feat: zstash ls --tars filters to tars containing matched files

docs/source/usage.rst

@@ -437,7 +437,10 @@ where
 * ``--hpss=<path to HPSS>`` specifies the destination path on the HPSS file system,
 * ``-l`` an optional argument to display more information.
 * ``--cache`` to use a cache other than the default of ``zstash``.
-* ``--tars`` to list the tars in addition to the files.
+* ``--tars`` to list the tars containing the matched files.
+  When combined with a ``[files]`` pattern, only the tars that hold those
+  matching files are shown — making it easy to identify the subset of tar
+  archives to download.
 * ``-v`` increases output verbosity.
 * ``[files]`` is a list of files to be listed (standard wildcards supported).
 
@@ -476,6 +479,26 @@ Below is an example of using ``ls`` to look at the tars in addition to the files
     Tars:
     000000.tar
 
+When combined with a file pattern, ``--tars`` shows only the tars that contain
+the matching files. This is useful for downloading a subset of the archive —
+for example, to find which tars hold files from 1850–1900 of a historical
+simulation::
+
+    $ zstash ls --hpss=hpss_archive --tars "*historical*185[0-9]*" "*historical*18[6-9][0-9]*" "*historical*190[0]*"
+
+    archive/run/historical.cam.h0.1850-01.nc
+    archive/run/historical.cam.h0.1851-06.nc
+    ...
+    archive/run/historical.cam.h0.1900-12.nc
+
+    Tars:
+    000000.tar
+    000001.tar
+    000007.tar
+
+You can then pass these tar names to ``zstash extract --tars`` to download only
+that subset.
+
 .. warning::
     Running ``zstash ls`` outside the source directory (the directory you're archiving)
     is not advised. ``zstash`` will only retrieve ``index.db`` from the HPSS archive

tests/unit/test_ls.py

@@ -0,0 +1,81 @@
+import datetime
+import os
+import sqlite3
+import tempfile
+import unittest
+from unittest.mock import patch
+
+from zstash.ls import ls_tars_database
+from zstash.settings import FilesRow
+
+
+class TestLsTarsDatabase(unittest.TestCase):
+    def setUp(self):
+        # Create an in-memory SQLite database with files and tars tables
+        self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
+        self.db_path = self.tmp.name
+        self.tmp.close()
+
+        con = sqlite3.connect(self.db_path)
+        cur = con.cursor()
+        cur.execute(
+            "CREATE TABLE tars (id INTEGER, name TEXT, size INTEGER, md5 TEXT)"
+        )
+        cur.execute(
+            "INSERT INTO tars VALUES (1, '000000.tar', 10240, 'abc123')"
+        )
+        cur.execute(
+            "INSERT INTO tars VALUES (2, '000001.tar', 20480, 'def456')"
+        )
+        cur.execute(
+            "INSERT INTO tars VALUES (3, '000002.tar', 30720, 'ghi789')"
+        )
+        con.commit()
+        con.close()
+
+    def tearDown(self):
+        os.unlink(self.db_path)
+
+    def _make_args(self, long: bool = False):
+        from unittest.mock import MagicMock
+
+        args = MagicMock()
+        args.long = long
+        return args
+
+    def test_ls_tars_database_filtered(self):
+        """Only tars containing matched files should be returned."""
+        args = self._make_args()
+        with patch("zstash.ls.get_db_filename", return_value=self.db_path):
+            result = ls_tars_database(args, "zstash", ["000000.tar", "000002.tar"])
+
+        names = [r.name for r in result]
+        self.assertIn("000000.tar", names)
+        self.assertIn("000002.tar", names)
+        self.assertNotIn("000001.tar", names)
+
+    def test_ls_tars_database_unfiltered(self):
+        """When tar_names is None, all tars should be returned."""
+        args = self._make_args()
+        with patch("zstash.ls.get_db_filename", return_value=self.db_path):
+            result = ls_tars_database(args, "zstash", None)
+
+        names = [r.name for r in result]
+        self.assertIn("000000.tar", names)
+        self.assertIn("000001.tar", names)
+        self.assertIn("000002.tar", names)
+
+    def test_tar_names_from_file_matches(self):
+        """The unique tar names extracted from FilesRow matches are correct."""
+        now = datetime.datetime(2024, 1, 1)
+        row1 = FilesRow((1, "archive/file_1850.nc", 100, now, "md5a", "000000.tar", 0))
+        row2 = FilesRow((2, "archive/file_1851.nc", 100, now, "md5b", "000000.tar", 512))
+        row3 = FilesRow((3, "archive/file_1900.nc", 100, now, "md5c", "000005.tar", 0))
+
+        matches = [row1, row2, row3]
+        tar_names = sorted(set(m.tar for m in matches))
+        self.assertEqual(tar_names, ["000000.tar", "000005.tar"])
+
+
+if __name__ == "__main__":
+    unittest.main()

zstash/ls.py

@@ -5,7 +5,7 @@
 import os
 import sqlite3
 import sys
-from typing import List, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 from .hpss import hpss_get
 from .settings import (
@@ -36,7 +36,8 @@ def ls():
     print_matches(args, matches)
 
     if args.tars:
-        tar_matches: List[TarsRow] = ls_tars_database(args, cache)
+        tar_names: List[str] = sorted(set(m.tar for m in matches))
+        tar_matches: List[TarsRow] = ls_tars_database(args, cache, tar_names)
         print_matches(args, tar_matches)
 
 
@@ -165,7 +166,9 @@ def ls_database(args: argparse.Namespace, cache: str) -> List[FilesRow]:
     return matches
 
 
-def ls_tars_database(args: argparse.Namespace, cache: str) -> List[TarsRow]:
+def ls_tars_database(
+    args: argparse.Namespace, cache: str, tar_names: Optional[List[str]] = None
+) -> List[TarsRow]:
     con: sqlite3.Connection = sqlite3.connect(
         get_db_filename(cache), detect_types=sqlite3.PARSE_DECLTYPES
     )
@@ -175,8 +178,14 @@ def ls_tars_database(args: argparse.Namespace, cache: str) -> List[TarsRow]:
         print("\ntars table does not exist")
         return []
 
-    # Find matching files
-    cur.execute("select * from tars")
+    # Find matching tars
+    if tar_names is not None:
+        placeholders = ",".join("?" * len(tar_names))
+        cur.execute(
+            "select * from tars where name in ({})".format(placeholders), tar_names
+        )
+    else:
+        cur.execute("select * from tars")
     matches_: List[TupleTarsRow] = cur.fetchall()
 
     # Remove duplicates