Black __init__.py

src/slurm_autoscale_tfe/__init__.py

@@ -26,12 +26,15 @@
 
 NODE_STATE_REGEX = re.compile(r"^NodeName=([a-z0-9-]*).*State=([A-Z_+]*).*$")
 DOWN_FLAG_SET = frozenset(["DOWN", "POWER_DOWN", "POWERED_DOWN", "POWERING_DOWN"])
-INSTANCE_TYPES = frozenset([
-    "aws_instance",
-    "azurerm_linux_virtual_machine",
-    "google_compute_instance",
-    "openstack_compute_instance_v2",
-])
+INSTANCE_TYPES = frozenset(
+    [
+        "aws_instance",
+        "azurerm_linux_virtual_machine",
+        "google_compute_instance",
+        "openstack_compute_instance_v2",
+    ]
+)
+
 
 class AutoscaleException(Exception):
     """Raised when something bad happened in autoscale main"""
@@ -66,7 +69,6 @@ def change_host_state(hostlist, state, reason=None):
         )
 
 
-
 def resume(hostlist=sys.argv[-1]):
     """Issue a request to Terraform cloud to power up the instances listed in
     hostlist.
@@ -92,6 +94,7 @@ def suspend(hostlist=sys.argv[-1]):
         return 1
     return 0
 
+
 def resume_fail(hostlist=sys.argv[-1]):
     """Issue a request to Terraform cloud to power down the instances listed in
     hostlist.
@@ -106,8 +109,7 @@ def resume_fail(hostlist=sys.argv[-1]):
 
 
 def connect_tfe_client():
-    """Return a TFE client object using environment variables for authentication
-    """
+    """Return a TFE client object using environment variables for authentication"""
     if "TFE_TOKEN" not in environ:
         raise AutoscaleException(
             f"{sys.argv[0]} requires environment variable TFE_TOKEN"
@@ -131,8 +133,7 @@ def connect_tfe_client():
 
 
 def get_pool_from_tfe(tfe_client):
-    """Retrieve id and content of POOL variable from Terraform cloud
-    """
+    """Retrieve id and content of POOL variable from Terraform cloud"""
     try:
         tfe_var = tfe_client.fetch_variable(POOL_VAR)
     except Timeout as exc:
@@ -149,6 +150,7 @@ def get_pool_from_tfe(tfe_client):
         return tfe_var["id"], frozenset(tfe_var["value"])
     return tfe_var["id"], frozenset()
 
+
 def get_instances_from_tfe(tfe_client):
     """Return all names of instances that are created in Terraform Cloud state."""
     try:
@@ -163,6 +165,7 @@ def get_instances_from_tfe(tfe_client):
             address_prefix = resource["attributes"]["address"].split("[")[0]
     return frozenset(instances), address_prefix
 
+
 def main(command, set_op, hostlist):
     """Issue a request to Terraform cloud to modify the pool variable of the
     workspace indicated by TFE_WORKSPACE environment variable using the operation
@@ -192,7 +195,7 @@ def main(command, set_op, hostlist):
     try:
         run_id = tfe_client.apply(
             f"Slurm {command.value} {hostlist}".strip(),
-            targets=[f'module.{address_prefix}["{hostname}"]' for hostname in hosts]
+            targets=[f'module.{address_prefix}["{hostname}"]' for hostname in hosts],
         )
     except Timeout as exc:
         raise AutoscaleException("Connection to Terraform cloud timeout (5s)") from exc