Issues/4314 change build tag import (#4329)

* Change where build tags dict is located to not require aws * Make typing proper * Update test imports * Move build_tag_dict to aws init * Import tag dict from the right place --------- Co-authored-by: Adam Novak <[email protected]>
DataBiosphere · Feb 4, 2023 · 9257042 · 9257042
1 parent d656409
commit 9257042
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 31 deletions.
diff --git a/src/toil/common.py b/src/toil/common.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import json
 import logging
 import os
 import pickle
@@ -36,6 +37,7 @@
                     ContextManager,
                     Dict,
                     List,
+                    MutableMapping,
                     Optional,
                     Set,
                     Tuple,
@@ -67,7 +69,7 @@
                       QueueSizeMessage,
                       gen_message_bus_path)
 from toil.fileStores import FileID
-from toil.lib.aws import zone_to_region
+from toil.lib.aws import zone_to_region, build_tag_dict_from_env
 from toil.lib.compatibility import deprecated
 from toil.lib.conversions import bytes2human, human2bytes
 from toil.lib.io import try_path

diff --git a/src/toil/jobStores/aws/jobStore.py b/src/toil/jobStores/aws/jobStore.py
@@ -35,6 +35,7 @@
 from botocore.exceptions import ClientError
 
 import toil.lib.encryption as encryption
+from toil.lib.aws import build_tag_dict_from_env
 from toil.fileStores import FileID
 from toil.jobStores.abstractJobStore import (AbstractJobStore,
                                              ConcurrentFileModificationException,
@@ -56,8 +57,7 @@
                                   ReadableTransformingPipe,
                                   WritablePipe)
 from toil.lib.aws.session import establish_boto3_session
-from toil.lib.aws.utils import (build_tag_dict_from_env,
-                                create_s3_bucket,
+from toil.lib.aws.utils import (create_s3_bucket,
                                 flatten_tags,
                                 get_bucket_region,
                                 get_object_for_url,

diff --git a/src/toil/lib/aws/__init__.py b/src/toil/lib/aws/__init__.py
@@ -25,6 +25,7 @@
                     Dict,
                     Iterable,
                     List,
+                    MutableMapping,
                     Optional,
                     TypeVar,
                     Union)
@@ -172,3 +173,23 @@ def running_on_ecs() -> bool:
     """
     # We only care about relatively current ECS
     return 'ECS_CONTAINER_METADATA_URI_V4' in os.environ
+
+def build_tag_dict_from_env(environment: MutableMapping[str, str] = os.environ) -> Dict[str, str]:
+    tags = dict()
+    owner_tag = environment.get('TOIL_OWNER_TAG')
+    if owner_tag:
+        tags.update({'Owner': owner_tag})
+
+    user_tags = environment.get('TOIL_AWS_TAGS')
+    if user_tags:
+        try:
+            json_user_tags = json.loads(user_tags)
+            if isinstance(json_user_tags, dict):
+                tags.update(json.loads(user_tags))
+            else:
+                logger.error('TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}')
+                exit(1)
+        except json.decoder.JSONDecodeError:
+            logger.error('TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}')
+            exit(1)
+    return tags
diff --git a/src/toil/lib/aws/utils.py b/src/toil/lib/aws/utils.py
@@ -28,7 +28,8 @@
                     Optional,
                     Set,
                     Union,
-                    cast)
+                    cast,
+                    MutableMapping)
 from urllib.parse import ParseResult
 
 from toil.lib.aws import session
@@ -41,10 +42,8 @@
                             retry)
 
 if sys.version_info >= (3, 8):
-    from typing import Literal, MutableMapping
+    from typing import Literal
 else:
-    from typing import MutableMapping
-
     from typing_extensions import Literal
 
 try:
@@ -403,28 +402,6 @@ def list_objects_for_url(url: ParseResult) -> List[str]:
         logger.debug('Found in %s items: %s', url, listing)
         return listing
 
-
-def build_tag_dict_from_env(environment: MutableMapping[str, str] = os.environ) -> Dict[str, str]:
-    tags = dict()
-    owner_tag = environment.get('TOIL_OWNER_TAG')
-    if owner_tag:
-        tags.update({'Owner': owner_tag})
-
-    user_tags = environment.get('TOIL_AWS_TAGS')
-    if user_tags:
-        try:
-            json_user_tags = json.loads(user_tags)
-            if isinstance(json_user_tags, dict):
-                tags.update(json.loads(user_tags))
-            else:
-                logger.error('TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}')
-                exit(1)
-        except json.decoder.JSONDecodeError:
-            logger.error('TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}')
-            exit(1)
-    return tags
-
-
 def flatten_tags(tags: Dict[str, str]) -> List[Dict[str, str]]:
     """
     Convert tags from a key to value dict into a list of 'Key': xxx, 'Value': xxx dicts.

diff --git a/src/toil/test/lib/aws/test_utils.py b/src/toil/test/lib/aws/test_utils.py
@@ -18,7 +18,7 @@
 
 import pytest
 
-from toil.lib.aws.utils import build_tag_dict_from_env
+from toil.lib.aws import build_tag_dict_from_env
 from toil.test import ToilTest
 
 logger = logging.getLogger(__name__)

diff --git a/src/toil/utils/toilLaunchCluster.py b/src/toil/utils/toilLaunchCluster.py
@@ -19,7 +19,7 @@
 
 from toil import applianceSelf
 from toil.common import parser_with_common_options
-from toil.lib.aws.utils import build_tag_dict_from_env
+from toil.lib.aws import build_tag_dict_from_env
 from toil.provisioners import (check_valid_node_types,
                                cluster_factory,
                                parse_node_types)