diff --git a/cosmicweb_music/cosmICweb.py b/cosmicweb_music/cosmICweb.py index 60a4d5a..26b68f6 100755 --- a/cosmicweb_music/cosmICweb.py +++ b/cosmicweb_music/cosmICweb.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import os import sys import tempfile @@ -26,7 +28,7 @@ EDITOR = os.environ.get("EDITOR", "vim") EDITOR_IS_VIM = EDITOR in {"vim", "nvim"} -def query_yes_no(question, default="yes"): +def query_yes_no(question: str, default="yes") -> bool: """Ask a yes/no question via raw_input() and return their answer. "question" is a string that is presented to the user. @@ -58,7 +60,7 @@ def query_yes_no(question, default="yes"): # Routines -def fetch_ellipsoids(url, api_token, attempts): +def fetch_ellipsoids(url: str, api_token: str, attempts: int) -> list[Ellipsoid]: for i in range(attempts): try: r = requests.get(url, headers={"Authorization": "Token " + api_token}) @@ -79,19 +81,21 @@ def fetch_ellipsoids(url, api_token, attempts): for e in content ] logging.error("Unable to download ellipsoids from {}".format(url)) - return None + return [] -def fetch_ellipsoid(url, api_token, traceback_radius, attempts=3): +def fetch_ellipsoid( + url: str, api_token: str, traceback_radius, attempts: int = 3 +) -> Ellipsoid | None: ellipsoids = fetch_ellipsoids(url, api_token, attempts) - if ellipsoids is not None: + if ellipsoids: return next( (e for e in ellipsoids if e.traceback_radius == traceback_radius), None ) return None -def fetch_downloadstore(cosmicweb_url, target): +def fetch_downloadstore(cosmicweb_url: str, target: str) -> DownloadConfig: try: r = requests.get(cosmicweb_url + "/api/music/store/" + target) # This will raise an error if not successful @@ -122,7 +126,9 @@ def fetch_downloadstore(cosmicweb_url, target): ) -def fetch_publication(cosmicweb_url, publication_name, traceback_radius): +def fetch_publication( + cosmicweb_url: str, publication_name: str, traceback_radius +) -> DownloadConfig: try: r = requests.get(cosmicweb_url + "/api/publications/" + publication_name) # This will raise an error if not successful @@ -155,7 +161,7 @@ def fetch_publication(cosmicweb_url, publication_name, traceback_radius): ) -def edit_template(template): +def edit_template(template: str) -> str: with tempfile.NamedTemporaryFile(suffix=".tmp.conf", mode="r+") as tf: tf.write(template) tf.flush() @@ -169,7 +175,7 @@ def edit_template(template): return template -def apply_config_parameter(config: str, parameters: dict[str, Any]): +def apply_config_parameter(config: str, parameters: dict[str, Any]) -> str: new_lines = [] for line in config.split("\n"): param = line.split("=")[0].strip() @@ -179,7 +185,7 @@ def apply_config_parameter(config: str, parameters: dict[str, Any]): return "\n".join(new_lines) -def music_config_to_template(config: DownloadConfig): +def music_config_to_template(config: DownloadConfig) -> str: music_config = config.MUSIC settings = config.settings # TODO: apply output configuration @@ -208,7 +214,7 @@ def compose_template( config: DownloadConfig, halo_name: str, halo_id: int, -): +) -> str: # TODO: add ellipsoid header (rtb, halo_name, etc) shape_0 = ", ".join(str(e) for e in ellipsoid.shape[0]) shape_1 = ", ".join(str(e) for e in ellipsoid.shape[1]) @@ -243,11 +249,11 @@ def write_music_file(output_file, music_config) -> None: f.write(music_config) -def call_music(): +def call_music() -> None: pass -def process_config(config: DownloadConfig, args: Args): +def process_config(config: DownloadConfig, args: Args) -> None: ellipsoids = [] for halo_name, url in zip(config.halo_names, config.halo_urls): logging.info("Fetching ellipsoids from halo " + halo_name) @@ -299,7 +305,7 @@ def process_config(config: DownloadConfig, args: Args): # TODO: Execute MUSIC? -def downloadstore_mode(args: Args, target: str): +def downloadstore_mode(args: Args, target: str) -> None: logging.info("Fetching download configuration from the cosmICweb server") config = fetch_downloadstore(args.url, target) if args.output_path == "./": @@ -309,7 +315,7 @@ def downloadstore_mode(args: Args, target: str): process_config(config, args) -def publication_mode(args: Args, publication_name: str, traceback_radius: int): +def publication_mode(args: Args, publication_name: str, traceback_radius: int) -> None: logging.info( "Fetching publication " + publication_name + " from the cosmICweb server" ) @@ -320,7 +326,7 @@ def publication_mode(args: Args, publication_name: str, traceback_radius: int): process_config(config, args) -def dir_path(p): +def dir_path(p: str) -> str: if os.path.isdir(p): return p else: @@ -373,54 +379,3 @@ def publication(ctx, publication_name, traceback_radius): if __name__ == "__main__": cli() - -# if __name__ == "__main__": -# parser = argparse.ArgumentParser() -# parser.add_argument( -# "--url", -# dest="cosmicweb_url", -# default=DEFAULT_URL, -# help="overwrite URL of the cosmicweb server", -# ) -# parser.add_argument( -# "--output-path", -# type=dir_path, -# default="./", -# help="Download target for IC files. If downloading publication, will create a subfolder with the " -# "name of the publication", -# ) -# parser.add_argument( -# "--common-directory", dest="create_subdirs", action="store_false" -# ) -# parser.add_argument( -# "--attempts", -# type=int, -# default=3, -# help="number of attempts to download ellipsoids", -# ) -# parser.add_argument("--verbose", action="store_true") - -# subparsers = parser.add_subparsers(dest="mode") -# # Downloading from publications -# publication_parser = subparsers.add_parser( -# "publication", help="download publications" -# ) -# publication_parser.add_argument("publication_name", help="name of the publication") -# publication_parser.add_argument( -# "--traceback_radius", type=int, choices=[1, 2, 4, 10], default=2, help="" -# ) -# # Downloading from download object -# download_parser = subparsers.add_parser("get") -# download_parser.add_argument("target") - -# args = parser.parse_args() - -# if args.verbose: -# logger.setLevel("DEBUG") - -# if args.mode == "get": -# downloadstore_mode(args) -# elif args.mode == "publication": -# publication_mode(args) -# else: -# raise NotImplementedError("unknown subparser") diff --git a/cosmicweb_music/data_types.py b/cosmicweb_music/data_types.py index e0afba7..0f07d1b 100644 --- a/cosmicweb_music/data_types.py +++ b/cosmicweb_music/data_types.py @@ -1,5 +1,6 @@ from __future__ import annotations +from datetime import datetime from typing import NamedTuple, Any, List, Dict, TypedDict