Sourcery refactored master branch

deployer/billing/importers.py

@@ -37,7 +37,7 @@ def build_gcp_query(cluster: dict, service_id=None):
         assert re.match(r"^[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}$", service_id, re.I)
         by_service = f'AND service.id = "{service_id}"'
 
-    query = f"""
+    return f"""
             SELECT
             invoice.month as month,
             project.id as project,
@@ -54,7 +54,6 @@ def build_gcp_query(cluster: dict, service_id=None):
             ORDER BY invoice.month ASC
             ;
             """
-    return query
 
 
 class BigqueryGCPBillingCostImporter:
@@ -135,8 +134,7 @@ def _run_query(self, start_month, end_month):
         start = start_month.timestamp()
         end = end_month.timestamp()
         step = "24h"
-        rows = p.query_range(prom_query, start, end, step)
-        return rows
+        return p.query_range(prom_query, start, end, step)
 
     def clean_query_dataframe(self, df):
         df = self.clean_namespace_labels(df)
@@ -165,7 +163,7 @@ def clean_namespace_labels(self, df):
     def combine_support(self, df):
         df["support_combined"] = 0.0
         if "support" in df:
-            df["support_combined"] = df["support_combined"] + df["support"]
+            df["support_combined"] += df["support"]
 
         if "kube-system" in df:
             df["support_combined"] = df["support_combined"] + df["kube-system"]
@@ -210,8 +208,7 @@ def get_dedicated_cluster_costs(cluster, start_month, end_month):
         pandas.DataFrame of costs.
     """
     bq_importer = BigqueryGCPBillingCostImporter(cluster)
-    result = bq_importer.get_costs(start_month, end_month)
-    return result
+    return bq_importer.get_costs(start_month, end_month)
 
 
 def get_shared_cluster_utilization(cluster, start_month, end_month):

deployer/cilogon_app.py

@@ -51,10 +51,7 @@ def build_request_headers(admin_id, admin_secret):
 def build_request_url(id=None):
     url = "https://cilogon.org/oauth2/oidc-cm"
 
-    if id is None:
-        return url
-
-    return str(URL(url).with_query({"client_id": id}))
+    return url if id is None else str(URL(url).with_query({"client_id": id}))
 
 
 def build_client_details(cluster_name, hub_name, callback_url):
@@ -125,10 +122,7 @@ def stored_client_id_same_with_cilogon_records(
         return False
 
     cilogon_client_id = cilogon_client.get("client_id", None)
-    if cilogon_client_id != stored_client_id:
-        return False
-
-    return True
+    return cilogon_client_id == stored_client_id
 
 
 def print_not_ok_request_message(response):
@@ -159,8 +153,7 @@ def create_client(
 
     # Check if there's a client id already stored in the config file for this hub
     if Path(config_filename).is_file():
-        client_id = load_client_id_from_file(config_filename)
-        if client_id:
+        if client_id := load_client_id_from_file(config_filename):
             print_colour(
                 f"Found existing CILogon client app in {config_filename}.", "yellow"
             )
@@ -294,23 +287,7 @@ def delete_client(admin_id, admin_secret, cluster_name, hub_name, client_id=None
         cluster_name, hub_name
     )
 
-    if not client_id:
-        if Path(config_filename).is_file():
-            client_id = load_client_id_from_file(config_filename)
-            # Nothing to do if no client has been found
-            if not client_id:
-                print_colour(
-                    "No `client_id` to delete was provided and couldn't find any in `config_filename`",
-                    "red",
-                )
-                return
-        else:
-            print_colour(
-                f"No `client_id` to delete was provided and couldn't find any {config_filename} file",
-                "red",
-            )
-            return
-    else:
+    if client_id:
         if not stored_client_id_same_with_cilogon_records(
             admin_id,
             admin_secret,
@@ -324,6 +301,21 @@ def delete_client(admin_id, admin_secret, cluster_name, hub_name, client_id=None
             )
             return
 
+    elif Path(config_filename).is_file():
+        client_id = load_client_id_from_file(config_filename)
+        # Nothing to do if no client has been found
+        if not client_id:
+            print_colour(
+                "No `client_id` to delete was provided and couldn't find any in `config_filename`",
+                "red",
+            )
+            return
+    else:
+        print_colour(
+            f"No `client_id` to delete was provided and couldn't find any {config_filename} file",
+            "red",
+        )
+        return
     print(f"Deleting the CILogon client details for {client_id}...")
     headers = build_request_headers(admin_id, admin_secret)
     response = requests.delete(build_request_url(client_id), headers=headers, timeout=5)

deployer/cluster.py

@@ -84,9 +84,7 @@ def deploy_support(self, cert_manager_version):
                 str(support_dir),
             ]
 
-            for values_file in values_files:
-                cmd.append(f"--values={values_file}")
-
+            cmd.extend(f"--values={values_file}" for values_file in values_files)
             print_colour(f"Running {' '.join([str(c) for c in cmd])}")
             subprocess.check_call(cmd)
 

deployer/deployer.py

@@ -132,7 +132,7 @@ def deploy_grafana_dashboards(
     # Add GRAFANA_TOKEN to the environment variables for
     # jupyterhub/grafana-dashboards' deploy.py script
     deploy_script_env = os.environ.copy()
-    deploy_script_env.update({"GRAFANA_TOKEN": grafana_token})
+    deploy_script_env["GRAFANA_TOKEN"] = grafana_token
     try:
         print_colour(f"Deploying grafana dashboards to {cluster_name}...")
         subprocess.check_call(
@@ -239,9 +239,9 @@ def generate_helm_upgrade_jobs(
             "reason_for_redeploy": "",
         }
 
-        # Check if this cluster file has been modified. If so, set boolean flags to True
-        intersection = set(changed_filepaths).intersection([str(cluster_file)])
-        if intersection:
+        if intersection := set(changed_filepaths).intersection(
+            [str(cluster_file)]
+        ):
             print_colour(
                 f"This cluster.yaml file has been modified. Generating jobs to upgrade all hubs and the support chart on THIS cluster: {cluster_name}"
             )
@@ -348,7 +348,7 @@ def run_hub_health_check(
     elif len(hub_indx) > 1:
         print_colour("ERROR: More than one hub with this name found!")
         sys.exit(1)
-    elif len(hub_indx) == 0:
+    elif not hub_indx:
         print_colour("ERROR: No hubs with this name found!")
         sys.exit(1)
 

deployer/file_acquisition.py

@@ -85,12 +85,10 @@ def build_absolute_path_to_hub_encrypted_config_file(cluster_name, hub_name):
         Note that file doesn't need to exist.
     """
     cluster_config_dir_path = find_absolute_path_to_cluster_file(cluster_name).parent
-    encrypted_file_path = cluster_config_dir_path.joinpath(
+    return cluster_config_dir_path.joinpath(
         f"enc-{hub_name}.secret.values.yaml"
     )
 
-    return encrypted_file_path
-
 
 def persist_config_in_encrypted_file(encrypted_file, new_config):
     """
@@ -196,10 +194,7 @@ def get_decrypted_file(original_filepath):
         with open(original_filepath) as f:
             # JSON files output by terraform have hard tabs, the *one* thing that is
             # valid JSON but not valid YAML
-            if ext.endswith("json"):
-                loader_func = json.load
-            else:
-                loader_func = yaml.load
+            loader_func = json.load if ext.endswith("json") else yaml.load
             try:
                 content = loader_func(f)
             except ScannerError:

deployer/generate/common.py

@@ -45,8 +45,8 @@ def generate_support_files(cluster_config_directory, vars):
     print_colour("Generating the prometheus credentials encrypted file...", "yellow")
     alphabet = string.ascii_letters + string.digits
     credentials = {
-        "username": "".join(secrets.choice(alphabet) for i in range(64)),
-        "password": "".join(secrets.choice(alphabet) for i in range(64)),
+        "username": "".join(secrets.choice(alphabet) for _ in range(64)),
+        "password": "".join(secrets.choice(alphabet) for _ in range(64)),
     }
     with open(
         REPO_ROOT / "config/clusters/templates/common/support.secret.values.yaml"

deployer/grafana/central_grafana.py

@@ -42,7 +42,7 @@ def build_datasource_details(cluster_name):
     # Get the credentials of this prometheus instance
     prometheus_creds = get_cluster_prometheus_creds(cluster_name)
 
-    datasource_details = {
+    return {
         "name": cluster_name,
         "type": "prometheus",
         "access": "proxy",
@@ -52,8 +52,6 @@ def build_datasource_details(cluster_name):
         "secureJsonData": {"basicAuthPassword": prometheus_creds["password"]},
     }
 
-    return datasource_details
-
 
 def build_datasource_request_headers(cluster_name):
     """
@@ -64,14 +62,12 @@ def build_datasource_request_headers(cluster_name):
     """
     token = get_grafana_token(cluster_name)
 
-    headers = {
+    return {
         "Accept": "application/json",
         "Content-Type": "application/json",
         "Authorization": f"Bearer {token}",
     }
 
-    return headers
-
 
 def get_clusters_used_as_datasources(cluster_name, datasource_endpoint):
     """
@@ -153,8 +149,6 @@ def update_central_grafana_datasources(
                     "yellow",
                 )
                 exceptions += 1
-                pass
-
     if exceptions:
         print_colour(
             f"Failed to add {exceptions} clusters as datasources. See errors above!",

deployer/grafana/grafana_utils.py

@@ -21,15 +21,14 @@ def get_grafana_url(cluster_name):
     with open(config_file) as f:
         support_config = yaml.load(f)
 
-    grafana_tls_config = (
+    if grafana_tls_config := (
         support_config.get("grafana", {}).get("ingress", {}).get("tls", [])
-    )
-    if not grafana_tls_config:
+    ):
+        # We only have one tls host right now. Modify this when things change.
+        return "https://" + grafana_tls_config[0]["hosts"][0]
+    else:
         raise ValueError(f"grafana.ingress.tls config for {cluster_name} missing!")
 
-    # We only have one tls host right now. Modify this when things change.
-    return "https://" + grafana_tls_config[0]["hosts"][0]
-
 
 def get_cluster_prometheus_address(cluster_name):
     """
@@ -56,21 +55,19 @@ def get_cluster_prometheus_address(cluster_name):
             f"`prometheusIngressAuthSecret` wasn't configured for {cluster_name}"
         )
 
-    tls_config = (
+    if tls_config := (
         support_config.get("prometheus", {})
         .get("server", {})
         .get("ingress", {})
         .get("tls", [])
-    )
-
-    if not tls_config:
+    ):
+        # We only have one tls host right now. Modify this when things change.
+        return tls_config[0]["hosts"][0]
+    else:
         raise ValueError(
             f"No tls config was found for the prometheus instance of {cluster_name}"
         )
 
-    # We only have one tls host right now. Modify this when things change.
-    return tls_config[0]["hosts"][0]
-
 
 def get_cluster_prometheus_creds(cluster_name):
     """

deployer/helm_upgrade_decision.py

@@ -151,11 +151,9 @@ def generate_hub_matrix_jobs(
                 cluster_file.parent.joinpath(values_file)
                 for values_file in hub.get("helm_chart_values_files", {})
             ]
-            # Establish if any of this hub's helm chart values files have been
-            # modified
-            intersection = added_or_modified_files.intersection(values_files)
-
-            if intersection:
+            if intersection := added_or_modified_files.intersection(
+                values_files
+            ):
                 # If at least one of the helm chart values files associated with
                 # this hub has been modified, add it to list of matrix jobs to be
                 # upgraded
@@ -228,9 +226,7 @@ def generate_support_matrix_jobs(
     # Empty list to store the matrix definitions in
     matrix_jobs = []
 
-    # Double-check that support is defined for this cluster.
-    support_config = cluster_config.get("support", {})
-    if support_config:
+    if support_config := cluster_config.get("support", {}):
         if upgrade_support_on_all_clusters or upgrade_support_on_this_cluster:
             # We know we're upgrading support on all clusters, so just add the cluster
             # name to the list of matrix jobs and move on
@@ -250,9 +246,9 @@ def generate_support_matrix_jobs(
                 cluster_file.parent.joinpath(values_file)
                 for values_file in support_config.get("helm_chart_values_files", {})
             ]
-            intersection = added_or_modified_files.intersection(values_files)
-
-            if intersection:
+            if intersection := added_or_modified_files.intersection(
+                values_files
+            ):
                 matrix_job = cluster_info.copy()
                 matrix_job["upgrade_support"] = True
                 matrix_job[
@@ -382,14 +378,11 @@ def ensure_support_staging_jobs_have_correct_keys(
     # upgrade_staging key present, even if we just set it to False
     for job in support_and_staging_matrix_jobs:
         if "upgrade_staging" not in job.keys():
-            # Get a list of prod hubs running on the same cluster this staging job will
-            # run on
-            hubs_on_this_cluster = [
+            if hubs_on_this_cluster := [
                 hub["hub_name"]
                 for hub in prod_hub_matrix_jobs
                 if hub["cluster_name"] == job["cluster_name"]
-            ]
-            if hubs_on_this_cluster:
+            ]:
                 # There are prod hubs on this cluster that require an upgrade, and so we
                 # also upgrade staging
                 job["upgrade_staging"] = True
@@ -431,9 +424,9 @@ def assign_staging_jobs_for_missing_clusters(
     support_staging_clusters = {
         job["cluster_name"] for job in support_and_staging_matrix_jobs
     }
-    missing_clusters = prod_hub_clusters.difference(support_staging_clusters)
-
-    if missing_clusters:
+    if missing_clusters := prod_hub_clusters.difference(
+        support_staging_clusters
+    ):
         # Generate support/staging jobs for clusters that don't have them but do have
         # prod hub jobs. We assume they are missing because neither the support chart
         # nor staging hub needed an upgrade. We set upgrade_support to False. However,

deployer/hub.py

@@ -57,9 +57,9 @@ def deploy(self, dask_gateway_version, debug, dry_run):
                 subprocess.check_call(["kubectl", "apply", "-f", manifest_url])
 
         with get_decrypted_files(
-            self.cluster.config_path.joinpath(p)
-            for p in self.spec["helm_chart_values_files"]
-        ) as values_files:
+                self.cluster.config_path.joinpath(p)
+                for p in self.spec["helm_chart_values_files"]
+            ) as values_files:
             cmd = [
                 "helm",
                 "upgrade",
@@ -78,9 +78,7 @@ def deploy(self, dask_gateway_version, debug, dry_run):
                 cmd.append("--debug")
 
             # Add on the values files
-            for values_file in values_files:
-                cmd.append(f"--values={values_file}")
-
+            cmd.extend(f"--values={values_file}" for values_file in values_files)
             # join method will fail on the PosixPath element if not transformed
             # into a string first
             print_colour(f"Running {' '.join([str(c) for c in cmd])}")