import SerializationContext from asdf.extension for version >3

BAMWelDX · Mar 22, 2024 · f9dca49 · f9dca49
1 parent e3fcbbb
commit f9dca49
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 4 deletions.
diff --git a/weldx/asdf/types.py b/weldx/asdf/types.py
@@ -3,7 +3,12 @@
 import functools
 import re
 
-from asdf.asdf import SerializationContext
+import asdf
+
+if asdf.__version__ >= "3.0.0":
+    from asdf.extension import SerializationContext
+else:
+    from asdf.asdf import SerializationContext
 from asdf.extension import Converter
 from asdf.versioning import AsdfSpec
 from boltons.iterutils import remap

diff --git a/weldx/asdf/util.py b/weldx/asdf/util.py
@@ -11,7 +11,11 @@
 
 import asdf
 import pint
-from asdf.asdf import SerializationContext
+
+if asdf.__version__ >= "3.0.0":
+    from asdf.extension import SerializationContext
+else:
+    from asdf.asdf import SerializationContext
 from asdf.config import AsdfConfig, get_config
 from asdf.extension import Extension
 from asdf.tagged import TaggedDict, TaggedList, TaggedString
@@ -495,7 +499,7 @@ def from_yaml_tree(
 
 def get_weldx_extension(ctx: SerializationContext | AsdfConfig) -> Extension:
     """Grab the weldx extension from list of current active extensions."""
-    if isinstance(ctx, asdf.asdf.SerializationContext):
+    if isinstance(ctx, SerializationContext):
         extensions = ctx.extension_manager.extensions
     elif isinstance(ctx, asdf.config.AsdfConfig):
         extensions = ctx.extensions

diff --git a/weldx/tags/base_types.py b/weldx/tags/base_types.py
@@ -2,7 +2,12 @@
 
 from uuid import UUID
 
-from asdf.asdf import SerializationContext
+import asdf
+
+if asdf.__version__ >= "3.0.0":
+    from asdf.extension import SerializationContext
+else:
+    from asdf.asdf import SerializationContext
 
 from weldx.asdf.types import WeldxConverter