Repository: popsim-consortium/stdpopsim (language: python)

Issue:
Refactor catalog
Closes #785

I haven't changed anything here, just rearranged the code. Also refactored the species testing code to use pytest.


Discussion / hints:
# [Codecov](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789?src=pr&el=h1) Report
> Merging [#789](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789?src=pr&el=desc) (df579d7) into [main](https://codecov.io/gh/popsim-consortium/stdpopsim/commit/76bf3bc06ba8478d9fac877a64b18de332c5e77e?el=desc) (76bf3bc) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/graphs/tree.svg?width=650&height=150&src=pr&token=AOkcaVfvhX)](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##             main     #789   +/-   ##
=======================================
  Coverage   99.48%   99.49%           
=======================================
  Files          36       54   +18     
  Lines        2337     2383   +46     
  Branches      279      279           
=======================================
+ Hits         2325     2371   +46     
  Misses          6        6           
  Partials        6        6           
```


| [Impacted Files](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [stdpopsim/catalog/AraTha/\_\_init\_\_.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQXJhVGhhL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/catalog/AraTha/demographic\_models.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQXJhVGhhL2RlbW9ncmFwaGljX21vZGVscy5weQ==) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/catalog/AraTha/genetic\_maps.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQXJhVGhhL2dlbmV0aWNfbWFwcy5weQ==) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/catalog/AraTha/species.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQXJhVGhhL3NwZWNpZXMucHk=) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/catalog/BosTau/\_\_init\_\_.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQm9zVGF1L19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/catalog/BosTau/demographic\_models.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQm9zVGF1L2RlbW9ncmFwaGljX21vZGVscy5weQ==) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/catalog/BosTau/species.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQm9zVGF1L3NwZWNpZXMucHk=) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/catalog/CanFam/\_\_init\_\_.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQ2FuRmFtL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/catalog/CanFam/genetic\_maps.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQ2FuRmFtL2dlbmV0aWNfbWFwcy5weQ==) | `100.00% <100.00%> (ø)` | |
| [stdpopsim/catalog/CanFam/species.py](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree#diff-c3RkcG9wc2ltL2NhdGFsb2cvQ2FuRmFtL3NwZWNpZXMucHk=) | `100.00% <100.00%> (ø)` | |
| ... and [33 more](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789?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/popsim-consortium/stdpopsim/pull/789?src=pr&el=footer). Last update [76bf3bc...df579d7](https://codecov.io/gh/popsim-consortium/stdpopsim/pull/789?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Fix patch:
diff --git a/maintenance/__main__.py b/maintenance/__main__.py
new file mode 100644
index 000000000..8a5da26c4
--- /dev/null
+++ b/maintenance/__main__.py
@@ -0,0 +1,6 @@
+# Entry point for running maintenance scripts.
+# Usage: python -m maintenance
+from . import main
+
+if __name__ == "__main__":
+    main.main()
diff --git a/maintenance/main.py b/maintenance/main.py
new file mode 100644
index 000000000..92e13e23b
--- /dev/null
+++ b/maintenance/main.py
@@ -0,0 +1,223 @@
+"""
+Command line interface for managing routine maintenance tasks.
+
+This interface is for stdpopsim developers only.
+"""
+import contextlib
+import shutil
+import string
+import pathlib
+
+import click
+import black
+
+import stdpopsim
+from . import ensembl
+
+species_template = string.Template(
+    """
+import stdpopsim
+
+from . import genome_data
+
+_chromosomes = []
+for name, data in genome_data.data["chromosomes"].items():
+    _chromosomes.append(
+        stdpopsim.Chromosome(
+            id=name,
+            length=data["length"],
+            synonyms=data["synonyms"],
+            mutation_rate=0,  # FILL ME IN
+            recombination_rate=0,  # FILL ME IN
+        )
+    )
+
+_genome = stdpopsim.Genome(
+    chromosomes=_chromosomes,
+    mutation_rate_citations=[], # ADD CITATIONS
+    recombination_rate_citations=[],  # ADD CITATIONS
+    assembly_name=genome_data.data["assembly_name"],
+    assembly_accession=genome_data.data["assembly_accession"],
+    assembly_citations=[],
+)
+
+_species = stdpopsim.Species(
+    id="$sps_id",
+    name="FIXME",
+    common_name="FIXME",
+    genome=_genome,
+    generation_time=0,  # FIXME
+    generation_time_citations=[],
+    population_size=0,  # FIXME
+    population_size_citations=[],
+)
+
+stdpopsim.register_species(_species)
+"""
+)
+
+species_test_template = string.Template(
+    """
+import stdpopsim
+from tests import test_species
+
+
+class TestSpecies(test_species.SpeciesTestBase):
+
+    species = stdpopsim.get_species("$sps_id")
+
+    # TODO specific tests for species data.
+
+class TestGenome(test_species.GenomeTestBase):
+
+    genome = stdpopsim.get_species("$sps_id").genome
+
+"""
+)
+
+
+def black_format(code):
+    return black.format_file_contents(code, fast=False, mode=black.FileMode())
+
+
+def ensembl_stdpopsim_id(ensembl_id):
+    tmp = ensembl_id.split("_")[:2]
+    sps_id = "".join([x[0:3].capitalize() for x in tmp])
+    if len(sps_id) != 6:
+        raise ValueError(f"Cannot extract six character id from {ensembl_id}")
+    return sps_id
+
+
+def catalog_path(sps_id):
+    return pathlib.Path(f"stdpopsim/catalog/{sps_id}")
+
+
+def write_catalog_stub(path, sps_id, ensembl_id):
+    """
+    Writes stub files to the catalog for a new species.
+    """
+    with open(path / "__init__.py", "w") as f:
+        print('"""', file=f)
+        print(f"Catalog definitions for {sps_id} (Ensembl ID='{ensembl_id}')", file=f)
+        print('"""', file=f)
+        print("from . import species  # noqa: F401", file=f)
+
+    species_code = species_template.substitute(sps_id=sps_id)
+    path = path / "species.py"
+    click.echo(f"Writing species definition stub to {path}")
+    with open(path, "w") as f:
+        f.write(black_format(species_code))
+
+    test_code = species_test_template.substitute(sps_id=sps_id)
+    test_path = pathlib.Path("tests") / f"test_{sps_id}.py"
+    click.echo(f"Writing species test stub to {test_path}")
+    with open(test_path, "w") as f:
+        f.write(black_format(test_code))
+
+
+class DataWriter:
+    """
+    Writes data obtained from upstream sources into the stdpopsim
+    package hierarchy.
+    """
+
+    def __init__(self):
+        self.ensembl_client = ensembl.EnsemblRestClient()
+
+    @contextlib.contextmanager
+    def write(self, path):
+        with open(path, "w") as f:
+            print("# File autogenerated from Ensembl REST API. Do not edit.", file=f)
+            yield f
+
+    def add_species(self, ensembl_id, force=False):
+        sps_id = ensembl_stdpopsim_id(ensembl_id)
+        click.echo(f"Adding new species {sps_id} for Ensembl ID {ensembl_id}")
+        root = catalog_path(sps_id)
+        if force:
+            shutil.rmtree(root, ignore_errors=True)
+        root.mkdir()
+        self.write_genome_data(ensembl_id)
+        write_catalog_stub(root, sps_id, ensembl_id)
+
+    def write_genome_data(self, ensembl_id):
+        sps_id = ensembl_stdpopsim_id(ensembl_id)
+        path = catalog_path(sps_id)
+        if not path.exists():
+            raise ValueError(
+                f"Directory {id} corresponding to {ensembl_id} does" + "not exist"
+            )
+        click.echo(f"Writing genome data for {sps_id} {ensembl_id}")
+        path = path / "genome_data.py"
+        data = self.ensembl_client.get_genome_data(ensembl_id)
+        code = f"data = {data}"
+
+        # Format the code with Black so we don't get noisy diffs
+        with self.write(path) as f:
+            f.write(black_format(code))
+
+    def write_ensembl_release(self):
+        release = self.ensembl_client.get_release()
+        click.echo(f"Using Ensembl release {release}")
+        path = pathlib.Path("stdpopsim/catalog/ensembl_info.py")
+        code = f"release = {release}"
+        with self.write(path) as f:
+            f.write(black_format(code))
+
+
+################
+# User interface
+################
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command()
+def list_species():
+    """
+    List species in stdpopsim with their Ensembl IDs.
+    """
+    click.echo("ID       Ensembl ID")
+    for species in stdpopsim.all_species():
+        click.echo(f"{species.id}   {species.ensembl_id}")
+
+
+@cli.command()
+@click.argument("species", nargs=-1)
+def update_genome_data(species):
+    """
+    Update the species genome data from Ensembl. Ensembl IDs for
+    species can optionally be provided, e.g.
+
+    update-genome-data homo_sapiens
+
+    will update the genome data for humans. Multiple species can
+    be specified. By default all species are updated.
+    """
+    if len(species) == 0:
+        embl_ids = [s.ensembl_id for s in stdpopsim.all_species()]
+    else:
+        embl_ids = [s.lower() for s in species]
+    writer = DataWriter()
+    for eid in embl_ids:
+        writer.write_genome_data(eid)
+    writer.write_ensembl_release()
+
+
+@cli.command()
+@click.argument("ensembl-id")
+@click.option("--force", is_flag=True)
+def add_species(ensembl_id, force):
+    """
+    Add a new species to the catalog using its ensembl ID.
+    """
+    writer = DataWriter()
+    writer.add_species(ensembl_id, force=force)
+    writer.write_ensembl_release()
+
+
+def main():
+    cli()
diff --git a/stdpopsim/__init__.py b/stdpopsim/__init__.py
index deed1a3b6..811b34d84 100644
--- a/stdpopsim/__init__.py
+++ b/stdpopsim/__init__.py
@@ -19,16 +19,9 @@
 from .engines import *  # NOQA
 from .warning_categories import *  # NOQA
 
-# Add imports for all defined species here.
-# We import these here to build the catalog, but the internal functions
+# We import catalog here, but the internal functions
 # defined are not part of the external API.
-from .catalog import AraTha  # NOQA
-from .catalog import BosTau  # NOQA
-from .catalog import CanFam  # NOQA
-from .catalog import DroMel  # NOQA
-from .catalog import EscCol  # NOQA
-from .catalog import HomSap  # NOQA
-from .catalog import PonAbe  # NOQA
+from .catalog import *  # NOQA
 
 from . import qc  # NOQA
 
diff --git a/stdpopsim/catalog/AraTha/__init__.py b/stdpopsim/catalog/AraTha/__init__.py
index 977e5cc0c..304063709 100644
--- a/stdpopsim/catalog/AraTha/__init__.py
+++ b/stdpopsim/catalog/AraTha/__init__.py
@@ -1,6 +1,6 @@
 """
 Catalog definitions for Arabidopsis Thaliana.
 """
-from . import species  # noqa 401
-from . import genetic_maps  # noqa 401
-from . import demographic_models  # noqa 401
+from . import species  # noqa: F401
+from . import genetic_maps  # noqa: F401
+from . import demographic_models  # noqa: F401
diff --git a/stdpopsim/catalog/BosTau/__init__.py b/stdpopsim/catalog/BosTau/__init__.py
index 30586d4d5..48461dd7f 100644
--- a/stdpopsim/catalog/BosTau/__init__.py
+++ b/stdpopsim/catalog/BosTau/__init__.py
@@ -1,5 +1,5 @@
 """
 Catalog definitions for Bos Taurus
 """
-from . import species  # noqa 401
-from . import demographic_models  # noqa 401
+from . import species  # noqa: F401
+from . import demographic_models  # noqa: F401
diff --git a/stdpopsim/catalog/CanFam/__init__.py b/stdpopsim/catalog/CanFam/__init__.py
index 268babc5a..1693f8185 100644
--- a/stdpopsim/catalog/CanFam/__init__.py
+++ b/stdpopsim/catalog/CanFam/__init__.py
@@ -1,5 +1,5 @@
 """
 Catalog data for Canis familiaris.
 """
-from . import species  # noqa 401
-from . import genetic_maps  # noqa 401
+from . import species  # noqa: F401
+from . import genetic_maps  # noqa: F401
diff --git a/stdpopsim/catalog/DroMel/__init__.py b/stdpopsim/catalog/DroMel/__init__.py
index 8ed4e71c3..62c2792f6 100644
--- a/stdpopsim/catalog/DroMel/__init__.py
+++ b/stdpopsim/catalog/DroMel/__init__.py
@@ -1,6 +1,6 @@
 """
 Catalog data for Drosophila Melanogaster
 """
-from . import species  # noqa 401
-from . import genetic_maps  # noqa 401
-from . import demographic_models  # noqa 401
+from . import species  # noqa: F401
+from . import genetic_maps  # noqa: F401
+from . import demographic_models  # noqa: F401
diff --git a/stdpopsim/catalog/DroMel/genome_data.py b/stdpopsim/catalog/DroMel/genome_data.py
index b5580ce14..eb3d7785b 100644
--- a/stdpopsim/catalog/DroMel/genome_data.py
+++ b/stdpopsim/catalog/DroMel/genome_data.py
@@ -1,7 +1,7 @@
 # File autogenerated from Ensembl REST API. Do not edit.
 data = {
     "assembly_accession": "GCA_000001215.4",
-    "assembly_name": "BDGP6.28",
+    "assembly_name": "BDGP6.32",
     "chromosomes": {
         "2L": {"length": 23513712, "synonyms": []},
         "2R": {"length": 25286936, "synonyms": []},
diff --git a/stdpopsim/catalog/EscCol/__init__.py b/stdpopsim/catalog/EscCol/__init__.py
index 8266c1ffa..2ffc87be0 100644
--- a/stdpopsim/catalog/EscCol/__init__.py
+++ b/stdpopsim/catalog/EscCol/__init__.py
@@ -1,4 +1,4 @@
 """
 Genome and demographic model definitions for Escherichia coli.
 """
-from . import species  # noqa 401
+from . import species  # noqa: F401
diff --git a/stdpopsim/catalog/EscCol/species.py b/stdpopsim/catalog/EscCol/species.py
index a95082c8c..91e14108d 100644
--- a/stdpopsim/catalog/EscCol/species.py
+++ b/stdpopsim/catalog/EscCol/species.py
@@ -60,7 +60,7 @@
     common_name="E. coli",
     # We use the K-12 strain, because the parameters we're using more
     # closely match this strain than the ensembl default (HUSEC2011).
-    ensembl_id="escherichia_coli_str_k_12_substr_mg1655",
+    ensembl_id="escherichia_coli_str_k_12_substr_mg1655_gca_000005845",
     genome=_genome,
     # E. coli K-12 strain MG1655 "doubling time during steady-state growth in
     # Luria-Bertani broth was 20 min".
diff --git a/stdpopsim/catalog/HomSap/__init__.py b/stdpopsim/catalog/HomSap/__init__.py
index f20451bd9..87c0c4484 100644
--- a/stdpopsim/catalog/HomSap/__init__.py
+++ b/stdpopsim/catalog/HomSap/__init__.py
@@ -1,7 +1,7 @@
 """
 Catalog definitions for Homo Sapiens
 """
-from . import species  # noqa 501
-from . import genetic_maps  # noqa 501
-from . import demographic_models  # noqa 501
-from . import annotations  # noqa 501
+from . import species  # noqa: F401
+from . import genetic_maps  # noqa: F401
+from . import demographic_models  # noqa: F401
+from . import annotations  # noqa: F401
diff --git a/stdpopsim/catalog/HomSap/species.py b/stdpopsim/catalog/HomSap/species.py
index 8289f2e19..23cade97e 100644
--- a/stdpopsim/catalog/HomSap/species.py
+++ b/stdpopsim/catalog/HomSap/species.py
@@ -2,13 +2,6 @@
 
 from . import genome_data
 
-###########################################################
-#
-# Genome definition
-#
-###########################################################
-
-
 # mean_recombination_rate was computed across all windows of the HapMapII_GRCh37 map
 _recombination_rate_data = {
     "1": 1.1485597641285933e-08,
diff --git a/stdpopsim/catalog/PonAbe/__init__.py b/stdpopsim/catalog/PonAbe/__init__.py
index fd0977028..33dd874ac 100644
--- a/stdpopsim/catalog/PonAbe/__init__.py
+++ b/stdpopsim/catalog/PonAbe/__init__.py
@@ -1,6 +1,6 @@
 """
 Catalog data for Pongo abelii
 """
-from . import species  # noqa 401
-from . import genetic_maps  # noqa 401
-from . import demographic_models  # noqa 401
+from . import species  # noqa: F401
+from . import genetic_maps  # noqa: F401
+from . import demographic_models  # noqa: F401
diff --git a/stdpopsim/catalog/__init__.py b/stdpopsim/catalog/__init__.py
index e69de29bb..a7d879f52 100644
--- a/stdpopsim/catalog/__init__.py
+++ b/stdpopsim/catalog/__init__.py
@@ -0,0 +1,10 @@
+import pathlib
+
+from . import ensembl_info  # noqa: F401
+
+# Import all species definitions in the catalog.
+__all__ = []
+for path in pathlib.Path(__path__[0]).glob("*"):
+    module_name = path.parts[-1]
+    if module_name[0].isupper():
+        __all__.append(module_name)
diff --git a/stdpopsim/catalog/ensembl_info.py b/stdpopsim/catalog/ensembl_info.py
index 1d96ff5d2..06d064d58 100644
--- a/stdpopsim/catalog/ensembl_info.py
+++ b/stdpopsim/catalog/ensembl_info.py
@@ -1,2 +1,2 @@
 # File autogenerated from Ensembl REST API. Do not edit.
-release = 101
+release = 103
diff --git a/update_ensembl_data.py b/update_ensembl_data.py
deleted file mode 100644
index 86d2874ca..000000000
--- a/update_ensembl_data.py
+++ /dev/null
@@ -1,95 +0,0 @@
-# Script used to update genome data from Ensembl.
-# TODO: this should be moved into the maintenance package,
-# ideally where we turn it into a CLI that we run with
-# python -m maintenance update-ensembl-data. That way we
-# can keep all the code for various scripts in one place,
-# and we don't clutter up the top-level directory with
-# scripts like this one.
-import contextlib
-
-import stdpopsim
-import maintenance
-import argparse
-import os
-
-
-class DataWriter:
-    """
-    Writes data obtained from upstream sources into the stdpopsim
-    package hierarchy.
-    """
-
-    def __init__(self):
-        self.ensembl_client = maintenance.EnsemblRestClient()
-
-    @contextlib.contextmanager
-    def write(self, path):
-        with open(path, "w") as f:
-            print("# File autogenerated from Ensembl REST API. Do not edit.", file=f)
-            yield f
-
-    def write_genome_data(self, ensembl_id):
-        tmp = ensembl_id.split("_")[:2]
-        id = "".join([x[0:3].capitalize() for x in tmp])
-        dir = path = f"stdpopsim/catalog/{id}"
-        if len(id) != 6:
-            raise ValueError(f"Cannot extract six character id from {ensembl_id}")
-        if not os.path.exists(dir):
-            raise ValueError(
-                f"Directory {id} corresponding to {ensembl_id} does" + "not exist"
-            )
-        print("Writing genome data for", id, ensembl_id)
-        path = f"{dir}/genome_data.py"
-        data = self.ensembl_client.get_genome_data(ensembl_id)
-        with self.write(path) as f:
-            print("data = ", data, file=f)
-
-    def write_ensembl_release(self):
-        release = self.ensembl_client.get_release()
-        print(f"Using Ensembl release {release}")
-        path = os.path.join("stdpopsim", "catalog", "ensembl_info.py")
-        with self.write(path) as f:
-            print("release = ", release, file=f)
-
-
-def main():
-    # Parser that controls which species should be updated
-    parser = argparse.ArgumentParser(
-        description="Query ensembl API for chromosome lengths"
-    )
-    parser.add_argument(
-        "species",
-        type=str,
-        nargs="*",
-        help="""Ensembl ids for species to update with underscores in place of
-                        spaces (e.g. homo_sapiens)""",
-    )
-    parser.add_argument(
-        "--list-species",
-        dest="list_species",
-        action="store_true",
-        help="list all species defined in stdpopsim",
-    )
-    args = parser.parse_args()
-
-    # Lists species and exits if user requested list of species
-    if args.list_species:
-        for species in stdpopsim.all_species():
-            print(species.ensembl_id)
-        return None
-
-    # Create a list of species ensembl ids based on user input
-    if len(args.species) == 0:
-        embl_ids = [s.ensembl_id for s in stdpopsim.all_species()]
-    else:
-        embl_ids = [s.lower() for s in args.species]
-
-    # Iterate over list of species ensembl ids and write genome data
-    writer = DataWriter()
-    writer.write_ensembl_release()
-    for eid in embl_ids:
-        writer.write_genome_data(eid)
-
-
-if __name__ == "__main__":
-    main()


Test patch:
diff --git a/tests/test_ensembl.py b/tests/test_ensembl.py
new file mode 100644
index 000000000..d28f72bed
--- /dev/null
+++ b/tests/test_ensembl.py
@@ -0,0 +1,7 @@
+import stdpopsim
+
+
+# Make sure we don't update the release without realising it.
+def test_version():
+    release = stdpopsim.catalog.ensembl_info.release
+    assert release == 103