style improvements, mostly

src/sunbather/__init__.py

@@ -6,45 +6,46 @@
 import shutil
 
 import sunbather.tools
+from sunbather.install_cloudy import GetCloudy
+
 
 def check_cloudy():
     """
     Checks if Cloudy executable exists, and if not, prompts to download and build it.
     """
     try:
-        CLOUDYVERSION = os.environ["CLOUDY_VERSION"]
+        cloudyversion = os.environ["CLOUDY_VERSION"]
     except KeyError:
-        CLOUDYVERSION = "23.01"
-    SUNBATHERPATH = os.path.dirname(
+        cloudyversion = "23.01"
+    sunbatherpath = os.path.dirname(
         os.path.abspath(__file__)
     )  # the absolute path where this code lives
     try:
         # the path where Cloudy is installed
-        CLOUDYPATH = os.environ["CLOUDY_PATH"]
-    except KeyError as exc:
-        CLOUDYPATH = f"{SUNBATHERPATH}/cloudy/c{CLOUDYVERSION}"
-    if not os.path.exists(f"{CLOUDYPATH}/source/cloudy.exe"):
+        cloudypath = os.environ["cloudy_path"]
+    except KeyError:
+        cloudypath = f"{sunbatherpath}/cloudy/c{cloudyversion}"
+    if not os.path.exists(f"{cloudypath}/source/cloudy.exe"):
         q = input(
             f"Cloudy not found and CLOUDY_PATH is not set. "
-            f"Do you want to install Cloudy {CLOUDYVERSION} now in the Sunbather path? "
+            f"Do you want to install Cloudy {cloudyversion} now in the sunbather path? "
             f"(y/n) "
         )
         while q.lower() not in ["y", "n"]:
             q = input("Please enter 'y' or 'n'")
         if q == "n":
             raise KeyError(
-                "Cloudy not found, and the environment variable 'CLOUDY_PATH' is not set. "
-                "Please set this variable in your .bashrc/.zshrc file "
+                "Cloudy not found, and the environment variable 'CLOUDY_PATH' is not "
+                "set. Please set this variable in your .bashrc/.zshrc file "
                 "to the path where the Cloudy installation is located. "
                 "Do not point it to the /source/ subfolder, but to the main folder."
-            ) from exc
-        from sunbather.install_cloudy import GetCloudy
-        INSTALLER = GetCloudy(version=CLOUDYVERSION)
-        INSTALLER.download()
-        INSTALLER.extract()
-        INSTALLER.compile()
-        INSTALLER.test()
-        INSTALLER.copy_data()
+            )
+        installer = GetCloudy(version=cloudyversion)
+        installer.download()
+        installer.extract()
+        installer.compile()
+        installer.test()
+        installer.copy_data()
 
 
 def make_workingdir():
@@ -64,7 +65,7 @@ def make_workingdir():
 
     sunbatherpath = f"{pathlib.Path(__file__).parent.resolve()}"
     shutil.copytree(
-        sunbatherpath + "/data/workingdir",
+        f"{sunbatherpath}/data/workingdir",
         workingdir,
     )
 
@@ -75,3 +76,5 @@ def firstrun():
     """
     check_cloudy()
     make_workingdir()
+
+    print("Sunbather is ready to go!")

src/sunbather/convergeT_parker.py

@@ -1,3 +1,7 @@
+"""
+ConvergeT_parker module of sunbather
+"""
+import sys
 import multiprocessing
 from shutil import copyfile
 import time
@@ -474,42 +478,44 @@ def run_g(
         Maximum number of iterations, by default 16.
     """
 
-    p = multiprocessing.Pool(cores)
-
-    pars = []
-    for Mdot in np.arange(
-        float(Mdot_l), float(Mdot_u) + 1e-6, float(Mdot_s)
-    ):  # 1e-6 so that upper bound is inclusive
-        for T in np.arange(int(T_l), int(T_u) + 1e-6, int(T_s)).astype(int):
-            pars.append(
-                (
-                    plname,
-                    Mdot,
-                    T,
-                    1,
-                    fc,
-                    workingdir,
-                    SEDname,
-                    overwrite,
-                    startT,
-                    pdir,
-                    zdict,
-                    altmax,
-                    save_sp,
-                    constantT,
-                    maxit,
+    with multiprocessing.Pool(processes=cores) as pool:
+        pars = []
+        for Mdot in np.arange(
+            float(Mdot_l), float(Mdot_u) + 1e-6, float(Mdot_s)
+        ):  # 1e-6 so that upper bound is inclusive
+            for T in np.arange(int(T_l), int(T_u) + 1e-6, int(T_s)).astype(int):
+                pars.append(
+                    (
+                        plname,
+                        Mdot,
+                        T,
+                        1,
+                        fc,
+                        workingdir,
+                        SEDname,
+                        overwrite,
+                        startT,
+                        pdir,
+                        zdict,
+                        altmax,
+                        save_sp,
+                        constantT,
+                        maxit,
+                    )
                 )
-            )
-
-    p.starmap(catch_errors_run_s, pars)
-    p.close()
-    p.join()
+        pool.starmap(catch_errors_run_s, pars)
+        pool.close()
+        pool.join()
 
 
-def main():
+def new_argument_parser():
     """
-    Main function
+    Creates a new argument parser.
     """
+    parser = argparse.ArgumentParser(
+        description="Runs the temperature convergence for 1D Parker profile(s).",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
 
     class OneOrThreeAction(argparse.Action):
         """
@@ -536,13 +542,6 @@ def __call__(self, parser, namespace, values, option_string=None):
                 key, val = value.split("=")
                 getattr(namespace, self.dest)[key] = float(val)
 
-    t0 = time.time()
-
-    parser = argparse.ArgumentParser(
-        description="Runs the temperature convergence for 1D Parker profile(s).",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-
     parser.add_argument(
         "-plname", required=True, help="planet name (must be in planets.txt)"
     )
@@ -691,7 +690,18 @@ def __call__(self, parser, namespace, values, option_string=None):
         ),
     )
 
-    args = parser.parse_args()
+    return parser
+
+
+def main(*args, **kwargs):
+    """
+    Main function
+    """
+
+    t0 = time.time()
+
+    parser = new_argument_parser()
+    args = parser.parse_args(*args, **kwargs)
 
     zdict = tools.get_zdict(z=args.z, zelem=args.zelem)
 
@@ -813,4 +823,4 @@ def __call__(self, parser, namespace, values, option_string=None):
 
 
 if __name__ == "__main__":
-    main()
+    main(sys.argv[1:])

src/sunbather/install_cloudy.py

@@ -1,3 +1,6 @@
+"""
+Functions to download and compile Cloudy
+"""
 import os
 import pathlib
 from urllib.error import HTTPError
@@ -37,7 +40,7 @@ def download(self):
             with urllib.request.urlopen(f"{self.url}{self.filename}") as g:
                 with open(self.filename, "b+w") as f:
                     f.write(g.read())
-        except HTTPError as exc:
+        except HTTPError:
             print(f"Could not download Cloudy from {self.url}{self.filename}...")
             return
         # Go to the v23 download page and download the "c23.01.tar.gz" file
@@ -56,26 +59,35 @@ def compile(self):
         Compiles Cloudy.
         """
         os.chdir(f"{self.cloudypath}/c{self.version}/source/")
-        subprocess.Popen(
+        with subprocess.Popen(
             [
                 "make",
             ]
-        ).wait()
+        ) as p:
+            p.wait()
 
     def test(self):
-        # Quickly test the Cloudy installation: in the source folder, run ./cloudy.exe, type "test" and hit return twice. It should print "Cloudy exited OK" at the end.
+        """
+        Quickly test the Cloudy installation: in the source folder, run
+        ./cloudy.exe, type "test" and hit return twice. It should print "Cloudy
+        exited OK" at the end.
+        """
         os.chdir(f"{self.cloudypath}/c{self.version}/source/")
         print(
             'Type "test" and hit return twice. '
             'It should print "Cloudy exited OK" at the end.'
         )
-        subprocess.Popen(
+        with subprocess.Popen(
             [
                 "./cloudy.exe",
             ]
-        ).wait()
+        ) as p:
+            p.wait()
 
     def copy_data(self):
+        """
+        Copy the stellar SEDs to the Cloudy data folder
+        """
         shutil.copytree(
             f"{self.sunbatherpath}/data/stellar_SEDs/",
             f"{self.cloudypath}/c{self.version}/data/SED/",