Add list-dates with counts CLI command

Author: jeromekelleher

sc2ts/cli.py

@@ -213,6 +213,27 @@ def initialise(ts, match_db, additional_problematic_sites, verbose, log_file):
     inference.MatchDb.initialise(match_db)
 
 
+@click.command()
+@click.argument("metadata", type=click.Path(exists=True, dir_okay=False))
+@click.option("--counts/--no-counts", default=False)
+@click.option("-v", "--verbose", count=True)
+@click.option("-l", "--log-file", default=None, type=click.Path(dir_okay=False))
+def list_dates(metadata, counts, verbose, log_file):
+    """
+    List the dates included in specified metadataDB
+    """
+    setup_logging(verbose, log_file)
+    with sc2ts.MetadataDb(metadata) as metadata_db:
+        counter = metadata_db.date_sample_counts()
+        if counts:
+            for k, v in counter.items():
+                print(k, v, sep="\t")
+
+        else:
+            for k in counter:
+                print(k)
+
+
 @click.command()
 @click.argument("alignments", type=click.Path(exists=True, dir_okay=False))
 @click.argument("metadata", type=click.Path(exists=True, dir_okay=False))
@@ -509,6 +530,7 @@ def cli():
 cli.add_command(export_metadata)
 
 cli.add_command(initialise)
+cli.add_command(list_dates)
 cli.add_command(daily_extend)
 cli.add_command(validate)
 cli.add_command(annotate_recombinants)

sc2ts/metadata.py

@@ -19,6 +19,7 @@ def __init__(self, path):
         uri = f"file:{path}"
         uri += "?mode=ro"
         self.uri = uri
+        self.path = path
         self.conn = sqlite3.connect(uri, uri=True)
         self.conn.row_factory = dict_factory
 
@@ -73,6 +74,14 @@ def get(self, date):
             for row in self.conn.execute(sql, [date]):
                 yield row
 
+    def date_sample_counts(self):
+        sql = "SELECT date, COUNT(*) FROM samples GROUP BY date ORDER BY date;"
+        counts = collections.Counter()
+        with self.conn:
+            for row in self.conn.execute(sql):
+                counts[row["date"]] = row["COUNT(*)"]
+        return counts
+
     def get_days(self, date=None):
         if date is None:
             date = "2000-01-01"

tests/test_cli.py

@@ -44,3 +44,67 @@ def test_additional_problematic_sites(self, tmp_path, additional):
         other_ts.tables.assert_equals(ts.tables)
         match_db = sc2ts.MatchDb(match_db_path)
         assert len(match_db) == 0
+
+
+class TestListDates:
+    def test_defaults(self, metadata_db):
+        runner = ct.CliRunner(mix_stderr=False)
+        result = runner.invoke(
+            cli.cli,
+            f"list-dates {metadata_db.path}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert result.stdout.splitlines() == [
+            "2020-01-01",
+            "2020-01-19",
+            "2020-01-24",
+            "2020-01-25",
+            "2020-01-28",
+            "2020-01-29",
+            "2020-01-30",
+            "2020-01-31",
+            "2020-02-01",
+            "2020-02-02",
+            "2020-02-03",
+            "2020-02-04",
+            "2020-02-05",
+            "2020-02-06",
+            "2020-02-07",
+            "2020-02-08",
+            "2020-02-09",
+            "2020-02-10",
+            "2020-02-11",
+            "2020-02-13",
+        ]
+
+    def test_counts(self, metadata_db):
+        runner = ct.CliRunner(mix_stderr=False)
+        result = runner.invoke(
+            cli.cli,
+            f"list-dates {metadata_db.path} --counts",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert result.stdout.splitlines() == [
+            "2020-01-01\t1",
+            "2020-01-19\t1",
+            "2020-01-24\t2",
+            "2020-01-25\t3",
+            "2020-01-28\t2",
+            "2020-01-29\t4",
+            "2020-01-30\t5",
+            "2020-01-31\t1",
+            "2020-02-01\t5",
+            "2020-02-02\t5",
+            "2020-02-03\t2",
+            "2020-02-04\t5",
+            "2020-02-05\t1",
+            "2020-02-06\t3",
+            "2020-02-07\t2",
+            "2020-02-08\t4",
+            "2020-02-09\t2",
+            "2020-02-10\t2",
+            "2020-02-11\t2",
+            "2020-02-13\t4",
+        ]