From d67b09f16d5badbc0a578250f470b6f0ff16ca84 Mon Sep 17 00:00:00 2001
From: Jon Massey <jon.massey@thedatalab.org>
Date: Thu, 28 Nov 2024 15:05:27 +0000
Subject: [PATCH] `codelists add` command

Allow users to add OpenCodelists codelists to the codelists.txt
file by URL.
---
 opensafely/codelists.py | 25 ++++++++++++++++++++++++-
 tests/test_codelists.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/opensafely/codelists.py b/opensafely/codelists.py
index f17ffe2..dac6581 100644
--- a/opensafely/codelists.py
+++ b/opensafely/codelists.py
@@ -29,6 +29,10 @@ def show_help(**kwargs):
         title="available commands", description="", metavar="COMMAND"
     )
 
+    parser_update = subparsers.add_parser(
+        "add", help="Add an OpenCodelists codelist to your project"
+    )
+
     parser_update = subparsers.add_parser(
         "update",
         help=(
@@ -63,6 +67,20 @@ def main():
     pass
 
 
+def add(codelist_url, codelists_dir=None):
+    if not codelists_dir:
+        codelists_dir = Path.cwd() / CODELISTS_DIR
+    codelists_file = get_codelist_file(codelists_dir)
+    if OPENCODELISTS_BASE_URL not in codelist_url:
+        exit_with_error(f"Unable to parse URL, {OPENCODELISTS_BASE_URL} not found")
+    line = codelist_url.replace(f"{OPENCODELISTS_BASE_URL}/codelist/", "")
+    with codelists_file.open("r+") as f:
+        if not f.readlines()[-1].endswith("\n"):
+            line = "\n" + line
+        f.write(line + "\n")
+    update(codelists_dir)
+
+
 def update(codelists_dir=None):
     if not codelists_dir:
         codelists_dir = Path.cwd() / CODELISTS_DIR
@@ -247,12 +265,17 @@ class Codelist:
     filename: Path
 
 
-def parse_codelist_file(codelists_dir):
+def get_codelist_file(codelists_dir):
     if not codelists_dir.exists() or not codelists_dir.is_dir():
         exit_with_error(f"No '{CODELISTS_DIR}' folder found")
     codelists_file = codelists_dir / CODELISTS_FILE
     if not codelists_file.exists():
         exit_with_error(f"No file found at '{CODELISTS_DIR}/{CODELISTS_FILE}'")
+    return codelists_file
+
+
+def parse_codelist_file(codelists_dir):
+    codelists_file = get_codelist_file(codelists_dir)
     codelists = []
     codelist_versions = {}
     for line in codelists_file.read_text().splitlines():
diff --git a/tests/test_codelists.py b/tests/test_codelists.py
index 3a32c34..cbe69fc 100644
--- a/tests/test_codelists.py
+++ b/tests/test_codelists.py
@@ -260,3 +260,31 @@ def test_codelists_check_with_upstream_changes_in_CI(
     os.chdir(codelists_path)
     # check doesn't fail in CI if there are upstream errors only
     assert codelists.check()
+
+
+def test_codelists_add(codelists_path, requests_mock):
+    codelists_path /= "codelists"
+    codelists_file = codelists_path / "codelists.txt"
+    prior_codelists = codelists_file.read_text()
+    for codelist in prior_codelists.split("\n"):
+        if codelist:
+            requests_mock.get(
+                f"https://www.opencodelists.org/codelist/{codelist.rstrip('/')}/download.csv",
+                text="foo",
+            )
+    requests_mock.get(
+        "https://www.opencodelists.org/"
+        "codelist/project123/codelist456/version1/download.csv",
+        text="foo",
+    )
+
+    codelists.add(
+        "https://www.opencodelists.org/codelist/project123/codelist456/version1",
+        codelists_path,
+    )
+
+    assert (
+        codelists_file.read_text()
+        == prior_codelists + "project123/codelist456/version1\n"
+    )
+    assert (codelists_path / "project123-codelist456.csv").read_text() == "foo"