From 1ee32928469096067f555a9d51f46b51d9595c3e Mon Sep 17 00:00:00 2001
From: Andrew Truong <itsandrewtruong@gmail.com>
Date: Wed, 21 Aug 2024 00:38:26 -0400
Subject: [PATCH] move to legacy: usage_analytics, weave_inspector

---
 weave/legacy/context_state.py         |  4 ++--
 weave/legacy/graph.py                 |  4 ++--
 weave/{ => legacy}/usage_analytics.py | 19 ++++++++++---------
 weave/{ => legacy}/weave_inspector.py |  4 ++--
 weave/query_api.py                    |  2 +-
 weave/weave_types.py                  |  4 ++--
 6 files changed, 19 insertions(+), 18 deletions(-)
 rename weave/{ => legacy}/usage_analytics.py (62%)
 rename weave/{ => legacy}/weave_inspector.py (99%)

diff --git a/weave/legacy/context_state.py b/weave/legacy/context_state.py
index 7ae2848dec37..15741b779878 100644
--- a/weave/legacy/context_state.py
+++ b/weave/legacy/context_state.py
@@ -218,11 +218,11 @@ def analytics_disabled():
         _analytics_enabled.reset(analytics_token)
 
 
-def analytics_enabled():
+def analytics_enabled() -> bool:
     return _analytics_enabled.get()
 
 
-def disable_analytics():
+def disable_analytics() -> contextvars.Token:
     return _analytics_enabled.set(False)
 
 
diff --git a/weave/legacy/graph.py b/weave/legacy/graph.py
index 51353cd5bb28..cb05793a81d8 100644
--- a/weave/legacy/graph.py
+++ b/weave/legacy/graph.py
@@ -6,7 +6,7 @@
 from weave.legacy import uris
 
 if typing.TYPE_CHECKING:
-    from weave import weave_inspector
+    from weave.legacy import weave_inspector
 
 
 T = typing.TypeVar("T")
@@ -39,7 +39,7 @@ def to_json(self) -> dict:
     def _inspect(self) -> "weave_inspector.NodeInspector":
         """Only intended to be used by developers to help debug the graph."""
         # Circular import, so we do it here.
-        from weave import weave_inspector
+        from weave.legacy import weave_inspector
 
         return weave_inspector.NodeInspector(self)
 
diff --git a/weave/usage_analytics.py b/weave/legacy/usage_analytics.py
similarity index 62%
rename from weave/usage_analytics.py
rename to weave/legacy/usage_analytics.py
index ce3d0aeb5713..a9fe6866a588 100644
--- a/weave/usage_analytics.py
+++ b/weave/legacy/usage_analytics.py
@@ -1,22 +1,23 @@
 import socket
+from typing import Optional
 
 import analytics
 
 from weave.legacy import context_state
 
-from . import environment
+from .. import environment
 
 analytics.write_key = "uJ8vZgKqTBVH6ZdhD4GZGZYsR7ucfJmb"
 
 
-def hostname():
+def hostname() -> str:
     return socket.gethostname()
 
 
 identify_called = False
 
 
-def _identify():
+def _identify() -> None:
     global identify_called
     if not identify_called:
         host = hostname()
@@ -24,14 +25,14 @@ def _identify():
     identify_called = True
 
 
-def analytics_enabled():
-    context_enabled = context_state.analytics_enabled()
-    env_enabled = environment.usage_analytics_enabled()
+def analytics_enabled() -> bool:
+    context_enabled: bool = context_state.analytics_enabled()
+    env_enabled: bool = environment.usage_analytics_enabled()
 
     return context_enabled and env_enabled
 
 
-def track(action: str, info=None):
+def track(action: str, info: Optional[dict] = None) -> None:
     if not analytics_enabled():
         return
     try:
@@ -43,9 +44,9 @@ def track(action: str, info=None):
         pass
 
 
-def use_called():
+def use_called() -> None:
     track("called use")
 
 
-def show_called(info=None):
+def show_called(info: Optional[dict] = None) -> None:
     track("called show", info)
diff --git a/weave/weave_inspector.py b/weave/legacy/weave_inspector.py
similarity index 99%
rename from weave/weave_inspector.py
rename to weave/legacy/weave_inspector.py
index b25a6c4ea21e..15814ada84ef 100644
--- a/weave/weave_inspector.py
+++ b/weave/legacy/weave_inspector.py
@@ -107,8 +107,8 @@
 
 from weave.legacy import graph
 
-from . import weave_types as types
-from .legacy.partial_object import PartialObjectType
+from .. import weave_types as types
+from .partial_object import PartialObjectType
 
 
 def _trimmed_string(s: str, max_len: int = 20) -> str:
diff --git a/weave/query_api.py b/weave/query_api.py
index 626d5345eba3..f457072c8e60 100644
--- a/weave/query_api.py
+++ b/weave/query_api.py
@@ -29,7 +29,7 @@
 from . import errors
 from weave.legacy.decorators import weave_class, mutation, type
 
-from . import usage_analytics
+from weave.legacy import usage_analytics
 from weave.legacy.context import (
     use_fixed_server_port,
     use_frontend_devmode,
diff --git a/weave/weave_types.py b/weave/weave_types.py
index e297bef18a39..9bc82f9039d1 100644
--- a/weave/weave_types.py
+++ b/weave/weave_types.py
@@ -20,7 +20,7 @@
     from weave.legacy import artifact_base
     from weave.legacy.artifact_fs import FilesystemArtifact
 
-    from . import weave_inspector
+    from .legacy import weave_inspector
 
 
 def to_weavejs_typekey(k: str) -> str:
@@ -456,7 +456,7 @@ def _make(cls, kwargs={}):
     def _inspect(self) -> "weave_inspector.TypeInspector":
         """Only intended to be used by developers to help debug the graph."""
         # Circular import, so we do it here.
-        from . import weave_inspector
+        from .legacy import weave_inspector
 
         return weave_inspector.TypeInspector(self)