From 4f6cd47411b96460495b1f7d6582860a4d655060 Mon Sep 17 00:00:00 2001
From: w-bonelli <wbonelli@ucar.edu>
Date: Thu, 31 Aug 2023 16:13:23 -0400
Subject: [PATCH] fix(generate_classes): use branch arg if provided (#1938)

* previously branch was ignored even if provided
* add repo param to download_dfn() function
* fix logging output
* fix test
---
 autotest/test_generate_classes.py   |  6 ++++--
 flopy/mf6/utils/generate_classes.py | 18 ++++++++----------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/autotest/test_generate_classes.py b/autotest/test_generate_classes.py
index fdd6b8ae80..cfbd4bd7a5 100644
--- a/autotest/test_generate_classes.py
+++ b/autotest/test_generate_classes.py
@@ -93,15 +93,17 @@ def test_generate_classes_from_github_refs(
     mod_file_times = [Path(mod_file).stat().st_mtime for mod_file in mod_files]
     pprint(mod_files)
 
-    # generate classes
+    # split ref into owner, repo, ref name
     spl = ref.split("/")
     owner = spl[0]
     repo = spl[1]
     ref = spl[2]
+
+    # generate classes
     pprint(
         virtualenv.run(
             "python -m flopy.mf6.utils.generate_classes "
-            f"--owner {owner} --repo {repo} ref {ref} --no-backup"
+            f"--owner {owner} --repo {repo} --ref {ref} --no-backup"
         )
     )
 
diff --git a/flopy/mf6/utils/generate_classes.py b/flopy/mf6/utils/generate_classes.py
index 602ed8fbb5..32c1d6978c 100644
--- a/flopy/mf6/utils/generate_classes.py
+++ b/flopy/mf6/utils/generate_classes.py
@@ -51,7 +51,7 @@ def list_files(pth, exts=["py"]):
             print(f"    {idx:5d} - {fn}")
 
 
-def download_dfn(owner, branch, new_dfn_pth):
+def download_dfn(owner, repo, ref, new_dfn_pth):
     try:
         from modflow_devtools.download import download_and_unzip
     except ImportError:
@@ -61,8 +61,7 @@ def download_dfn(owner, branch, new_dfn_pth):
             "     pip install modflow-devtools"
         )
 
-    mf6url = "https://github.com/{}/modflow6/archive/{}.zip"
-    mf6url = mf6url.format(owner, branch)
+    mf6url = f"https://github.com/{owner}/{repo}/archive/{ref}.zip"
     print(f"  Downloading MODFLOW 6 repository from {mf6url}")
     with tempfile.TemporaryDirectory() as tmpdirname:
         dl_path = download_and_unzip(mf6url, tmpdirname, verbose=True)
@@ -159,20 +158,19 @@ def generate_classes(
     # download the dfn files and put them in flopy.mf6.data or update using
     # user provided dfnpath
     if dfnpath is None:
-        print(
-            f"  Updating the MODFLOW 6 classes using {owner}/{repo}/{branch}"
-        )
         timestr = time.strftime("%Y%m%d-%H%M%S")
         new_dfn_pth = os.path.join(flopypth, "mf6", "data", f"dfn_{timestr}")
 
         # branch deprecated 3.5.0
-        if not ref:
-            if not branch:
-                raise ValueError("branch or ref must be provided")
+        if not ref and not branch:
+            raise ValueError("branch or ref must be provided")
+        if branch:
             warn("branch is deprecated, use ref instead", DeprecationWarning)
             ref = branch
 
-        download_dfn(owner, ref, new_dfn_pth)
+        print(f"  Updating the MODFLOW 6 classes using {owner}/{repo}/{ref}")
+
+        download_dfn(owner, repo, ref, new_dfn_pth)
     else:
         print(f"  Updating the MODFLOW 6 classes using {dfnpath}")
         assert os.path.isdir(dfnpath)