move to trace: ref_util

wandb · Aug 28, 2024 · 38e9c84 · 38e9c84
1 parent f8759ac
commit 38e9c84
Show file tree

Hide file tree

Showing 6 changed files with 10 additions and 11 deletions.
diff --git a/weave/legacy/arrow/list_.py b/weave/legacy/arrow/list_.py
@@ -12,7 +12,6 @@
 
 from weave import (
     errors,
-    ref_util,
     weave_internal,
 )
 from weave import weave_types as types
@@ -40,6 +39,7 @@
     tag_store,
     tagged_value_type,
 )
+from weave.trace import ref_util
 
 
 def reverse_dict(d: dict) -> dict:

diff --git a/weave/legacy/artifact_fs.py b/weave/legacy/artifact_fs.py
@@ -6,10 +6,11 @@
 import os
 import typing
 
-from weave import errors, ref_util
+from weave import errors
 from weave import weave_types as types
 from weave.legacy import artifact_base, file_base, object_context, ref_base, uris
 from weave.legacy.language_features.tagging import tag_store
+from weave.trace import ref_util
 
 if typing.TYPE_CHECKING:
     from weave.legacy import graph

diff --git a/weave/legacy/box.py b/weave/legacy/box.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from weave import ref_util
+from weave.trace import ref_util
 from weave.legacy import context_state
 
 

diff --git a/weave/legacy/object_type_ref_util.py b/weave/legacy/object_type_ref_util.py
@@ -1,6 +1,6 @@
 import typing
 
-from weave import ref_util
+from weave.trace import ref_util
 from weave.legacy import context_state
 
 

diff --git a/weave/tests/legacy/test_refs.py b/weave/tests/legacy/test_refs.py
@@ -4,13 +4,14 @@
 from weave.flow.obj import Object
 from weave.legacy import artifact_local
 from weave.legacy import ops_arrow as arrow
+from weave.trace import ref_util
 from weave.trace_server.refs_internal import (
     DICT_KEY_EDGE_NAME,
     LIST_INDEX_EDGE_NAME,
     OBJECT_ATTR_EDGE_NAME,
 )
 
-from ... import ref_util, storage
+from ... import storage
 
 
 def test_laref_artifact_version_1():

diff --git a/weave/ref_util.py → weave/trace/ref_util.py b/weave/ref_util.py → weave/trace/ref_util.py
@@ -3,8 +3,7 @@
 from urllib import parse
 
 from weave.legacy import box
-
-from .trace_server import refs_internal
+from weave.trace_server import refs_internal
 
 DICT_KEY_EDGE_NAME = refs_internal.DICT_KEY_EDGE_NAME
 LIST_INDEX_EDGE_NAME = refs_internal.LIST_INDEX_EDGE_NAME
@@ -23,9 +22,7 @@ def parse_local_ref_str(s: str) -> typing.Tuple[str, typing.Optional[list[str]]]
 def val_with_relative_ref(
     parent_object: typing.Any, child_object: typing.Any, ref_extra_parts: list[str]
 ) -> typing.Any:
-    from weave.legacy import context_state
-
-    from .legacy import ref_base
+    from weave.legacy import context_state, ref_base
 
     # If we already have a ref, resolve it
     if isinstance(child_object, ref_base.Ref):
@@ -34,7 +31,7 @@ def val_with_relative_ref(
     # Only do this if ref_tracking_enabled right now. I just want to
     # avoid introducing new behavior into W&B prod for the moment.
     if context_state.ref_tracking_enabled():
-        from . import storage
+        from .. import storage
 
         child_ref = storage.get_ref(child_object)
         parent_ref = ref_base.get_ref(parent_object)