Refactor: Use random.choices (#33631)

apache · Aug 24, 2023 · ba0bab0 · ba0bab0
1 parent 66d0222
commit ba0bab0
Show file tree

Hide file tree

Showing 8 changed files with 10 additions and 12 deletions.
diff --git a/airflow/auth/managers/fab/cli_commands/user_command.py b/airflow/auth/managers/fab/cli_commands/user_command.py
@@ -70,7 +70,7 @@ def users_create(args):
             valid_roles = appbuilder.sm.get_all_roles()
             raise SystemExit(f"{args.role} is not a valid role. Valid roles are: {valid_roles}")
         if args.use_random_password:
-            password = "".join(random.choice(string.printable) for _ in range(16))
+            password = "".join(random.choices(string.printable, k=16))
         elif args.password:
             password = args.password
         else:

diff --git a/airflow/cli/commands/standalone_command.py b/airflow/cli/commands/standalone_command.py
@@ -193,9 +193,8 @@ def initialize_database(self):
             self.print_output("standalone", "Creating admin user")
             role = appbuilder.sm.find_role("Admin")
             assert role is not None
-            password = "".join(
-                random.choice("abcdefghkmnpqrstuvwxyzABCDEFGHKMNPQRSTUVWXYZ23456789") for i in range(16)
-            )
+            # password does not contain visually similar characters: ijlIJL1oO0
+            password = "".join(random.choices("abcdefghkmnpqrstuvwxyzABCDEFGHKMNPQRSTUVWXYZ23456789", k=16))
             with open(password_path, "w") as file:
                 file.write(password)
             make_group_other_inaccessible(password_path)

diff --git a/airflow/providers/google/cloud/hooks/cloud_sql.py b/airflow/providers/google/cloud/hooks/cloud_sql.py
@@ -887,7 +887,7 @@ def _generate_unique_path() -> str:
         random.seed()
         while True:
             candidate = os.path.join(
-                gettempdir(), "".join(random.choice(string.ascii_lowercase + string.digits) for _ in range(8))
+                gettempdir(), "".join(random.choices(string.ascii_lowercase + string.digits, k=8))
             )
             if not os.path.exists(candidate):
                 return candidate

diff --git a/airflow/providers/google/cloud/hooks/mlengine.py b/airflow/providers/google/cloud/hooks/mlengine.py
@@ -65,13 +65,13 @@ def _poll_with_exponential_delay(
                 log.info("Operation is done: %s", response)
                 return response
 
-            time.sleep((2**i) + (random.randint(0, 1000) / 1000))
+            time.sleep((2**i) + random.random())
         except HttpError as e:
             if e.resp.status != 429:
                 log.info("Something went wrong. Not retrying: %s", format(e))
                 raise
             else:
-                time.sleep((2**i) + (random.randint(0, 1000) / 1000))
+                time.sleep((2**i) + random.random())
 
     raise ValueError(f"Connection could not be established after {max_n} retries.")
 

diff --git a/tests/providers/elasticsearch/log/elasticmock/utilities/__init__.py b/tests/providers/elasticsearch/log/elasticmock/utilities/__init__.py
@@ -54,7 +54,7 @@
 
 def get_random_id(size=DEFAULT_ELASTICSEARCH_ID_SIZE):
     """Returns random if for elasticsearch"""
-    return "".join(random.choice(CHARSET_FOR_ELASTICSEARCH_ID) for _ in range(size))
+    return "".join(random.choices(CHARSET_FOR_ELASTICSEARCH_ID, k=size))
 
 
 def query_params(*es_query_params, **kwargs):

diff --git a/tests/providers/ssh/hooks/test_ssh.py b/tests/providers/ssh/hooks/test_ssh.py
@@ -77,7 +77,7 @@ def generate_host_key(pkey: paramiko.PKey):
 TEST_CMD_TIMEOUT_NOT_SET = "NOT SET"
 TEST_CMD_TIMEOUT_EXTRA = 15
 
-PASSPHRASE = "".join(random.choice(string.ascii_letters) for i in range(10))
+PASSPHRASE = "".join(random.choices(string.ascii_letters, k=10))
 TEST_ENCRYPTED_PRIVATE_KEY = generate_key_string(pkey=TEST_PKEY, passphrase=PASSPHRASE)
 
 TEST_DISABLED_ALGORITHMS = {"pubkeys": ["rsa-sha2-256", "rsa-sha2-512"]}

diff --git a/tests/system/providers/apache/kafka/example_dag_event_listener.py b/tests/system/providers/apache/kafka/example_dag_event_listener.py
@@ -69,8 +69,7 @@ def _producer_function():
 
 
 def _generate_uuid():
-    letters = string.ascii_lowercase
-    return "".join(random.choice(letters) for i in range(6))
+    return "".join(random.choices(string.ascii_lowercase, k=6))
 
 
 with DAG(

diff --git a/tests/test_utils/azure_system_helpers.py b/tests/test_utils/azure_system_helpers.py
@@ -146,7 +146,7 @@ def prepare_share(cls, share_name: str, azure_fileshare_conn_id: str, file_name:
         cls.create_directory(
             share_name=share_name, azure_fileshare_conn_id=azure_fileshare_conn_id, directory=directory
         )
-        string_data = "".join(random.choice(string.ascii_letters) for _ in range(1024))
+        string_data = "".join(random.choices(string.ascii_letters, k=1024))
         cls.upload_file_from_string(
             string_data=string_data,
             share_name=share_name,