From 38e9c84d57703219dcc2b15fe397ce2aa826229b Mon Sep 17 00:00:00 2001
From: Andrew Truong <itsandrewtruong@gmail.com>
Date: Fri, 23 Aug 2024 08:48:56 -0400
Subject: [PATCH] move to trace: ref_util

---
 weave/legacy/arrow/list_.py          | 2 +-
 weave/legacy/artifact_fs.py          | 3 ++-
 weave/legacy/box.py                  | 2 +-
 weave/legacy/object_type_ref_util.py | 2 +-
 weave/tests/legacy/test_refs.py      | 3 ++-
 weave/{ => trace}/ref_util.py        | 9 +++------
 6 files changed, 10 insertions(+), 11 deletions(-)
 rename weave/{ => trace}/ref_util.py (95%)

diff --git a/weave/legacy/arrow/list_.py b/weave/legacy/arrow/list_.py
index f8a7f09bbb3..7cc6e84ddba 100644
--- 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
index f09e110e612..d5c71f3c643 100644
--- 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
index f089e03acdc..b957fc6ba3f 100644
--- 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
index 40e70fff6fd..53787c4c454 100644
--- 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
index a1cfd4ccca7..ac9923eaaad 100644
--- 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 b/weave/trace/ref_util.py
similarity index 95%
rename from weave/ref_util.py
rename to weave/trace/ref_util.py
index 968a584ca61..93ec19decd7 100644
--- a/weave/ref_util.py
+++ b/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)