From fdffde00117c164d93148e9a9aa7d116bd1f20fe Mon Sep 17 00:00:00 2001
From: Andrew Truong <itsandrewtruong@gmail.com>
Date: Thu, 22 Aug 2024 22:50:02 -0400
Subject: [PATCH] move to legacy: registry_mem

---
 weave/legacy/codify.py                         |  4 ++--
 weave/legacy/compile.py                        |  2 +-
 weave/legacy/compile_domain.py                 |  4 ++--
 weave/legacy/decorator_class.py                |  4 ++--
 weave/legacy/decorator_op.py                   |  3 +--
 weave/legacy/derive_op.py                      |  2 +-
 weave/legacy/dispatch.py                       |  4 ++--
 weave/legacy/ecosystem/root.py                 |  4 ++--
 weave/legacy/ecosystem/wandb/wandb_objs.py     |  2 +-
 weave/legacy/execute.py                        |  2 +-
 weave/legacy/execute_fast.py                   |  2 +-
 .../tagging/process_opdef_output_type.py       |  2 +-
 weave/legacy/op_def_type.py                    |  4 ++--
 weave/legacy/ops_arrow/vectorize.py            |  3 +--
 .../run_history/history_op_common.py           |  2 +-
 weave/legacy/ops_primitives/weave_api.py       |  2 +-
 weave/legacy/panels_py/generator_templates.py  |  4 ++--
 weave/legacy/propagate_gql_keys.py             |  4 ++--
 weave/{ => legacy}/registry_mem.py             | 18 +++++++++---------
 weave/legacy/stitch.py                         |  4 ++--
 weave/tests/legacy/test_derive_op.py           |  2 +-
 weave/tests/legacy/test_mappability.py         |  2 +-
 weave/tests/legacy/test_op_behaviors.py        |  3 ++-
 weave/tests/legacy/test_op_coverage.py         |  2 +-
 weave/tests/legacy/test_op_serialization.py    |  3 ++-
 weave/tests/legacy/test_serialize.py           |  4 ++--
 weave/tests/legacy/test_wb_domain_ops.py       |  2 +-
 weave/weave_internal.py                        |  2 +-
 weave/weave_server.py                          | 10 ++++++++--
 29 files changed, 56 insertions(+), 50 deletions(-)
 rename weave/{ => legacy}/registry_mem.py (92%)

diff --git a/weave/legacy/codify.py b/weave/legacy/codify.py
index 04fac73cb3c..0395775e034 100644
--- a/weave/legacy/codify.py
+++ b/weave/legacy/codify.py
@@ -6,8 +6,8 @@
 
 import black
 
-from weave import registry_mem, storage, weave_types
-from weave.legacy import graph
+from weave import storage, weave_types
+from weave.legacy import graph, registry_mem
 
 from . import codifiable_value_mixin
 
diff --git a/weave/legacy/compile.py b/weave/legacy/compile.py
index cbe4293e893..30cbdf4d98f 100644
--- a/weave/legacy/compile.py
+++ b/weave/legacy/compile.py
@@ -8,7 +8,6 @@
 from weave import (
     engine_trace,
     errors,
-    registry_mem,
     weave_internal,
 )
 from weave import weave_types as types
@@ -26,6 +25,7 @@
     op_args,
     partial_object,
     propagate_gql_keys,
+    registry_mem,
     serialize,
     stitch,
     value_or_error,
diff --git a/weave/legacy/compile_domain.py b/weave/legacy/compile_domain.py
index 7e04250e4f9..2a0f686d676 100644
--- a/weave/legacy/compile_domain.py
+++ b/weave/legacy/compile_domain.py
@@ -2,9 +2,9 @@
 
 import graphql
 
-from weave import errors, registry_mem
+from weave import errors
 from weave import weave_types as types
-from weave.legacy import gql_op_plugin, gql_to_weave, graph, op_args, stitch
+from weave.legacy import gql_op_plugin, gql_to_weave, graph, op_args, stitch, registry_mem
 from weave.legacy.input_provider import InputAndStitchProvider
 
 if typing.TYPE_CHECKING:
diff --git a/weave/legacy/decorator_class.py b/weave/legacy/decorator_class.py
index 2e42b15b1be..99fadbcb16e 100644
--- a/weave/legacy/decorator_class.py
+++ b/weave/legacy/decorator_class.py
@@ -1,9 +1,9 @@
 import inspect
 import typing
 
-from weave import errors, registry_mem
+from weave import errors
 from weave import weave_types as types
-from weave.legacy import context_state, derive_op, op_def
+from weave.legacy import context_state, derive_op, op_def, registry_mem
 
 # Contrary to the way it is read, the weave.class() decorator runs AFTER the
 # inner methods are defined. Therefore, this function runs after the ops are
diff --git a/weave/legacy/decorator_op.py b/weave/legacy/decorator_op.py
index c7f0dc8b6cb..e193e4b1e38 100644
--- a/weave/legacy/decorator_op.py
+++ b/weave/legacy/decorator_op.py
@@ -4,9 +4,8 @@
 
 from typing_extensions import ParamSpec
 
-from weave import registry_mem
 from weave import weave_types as types
-from weave.legacy import context_state, derive_op, op_args, op_def, pyfunc_type_util
+from weave.legacy import context_state, derive_op, op_args, op_def, pyfunc_type_util, registry_mem
 
 if typing.TYPE_CHECKING:
     from weave.legacy.gql_op_plugin import GqlOpPlugin
diff --git a/weave/legacy/derive_op.py b/weave/legacy/derive_op.py
index 13c0dcd35de..337bb66cccd 100644
--- a/weave/legacy/derive_op.py
+++ b/weave/legacy/derive_op.py
@@ -9,12 +9,12 @@
 from weave import (
     errors,
     parallelism,
-    registry_mem,
     storage,
     weave_internal,
 )
 from weave import weave_types as types
 from weave.legacy import (
+    registry_mem,
     box,
     context_state,
     execute_fast,
diff --git a/weave/legacy/dispatch.py b/weave/legacy/dispatch.py
index b0af4369b9f..4eaf2ca215f 100644
--- a/weave/legacy/dispatch.py
+++ b/weave/legacy/dispatch.py
@@ -5,9 +5,9 @@
 import typing
 from dataclasses import dataclass
 
-from weave import errors, registry_mem, util
+from weave import errors, util
 from weave import weave_types as types
-from weave.legacy import graph, memo, op_args, op_def, pyfunc_type_util
+from weave.legacy import graph, memo, op_args, op_def, pyfunc_type_util, registry_mem
 from weave.legacy.language_features.tagging.is_tag_getter import is_tag_getter
 from weave.legacy.language_features.tagging.tagged_value_type import TaggedValueType
 
diff --git a/weave/legacy/ecosystem/root.py b/weave/legacy/ecosystem/root.py
index 0757c9b03d3..38cd09d3e7e 100644
--- a/weave/legacy/ecosystem/root.py
+++ b/weave/legacy/ecosystem/root.py
@@ -22,7 +22,7 @@
 """
 )
 
-from weave import registry_mem
+from weave.legacy import registry_mem
 
 op_org_name = registry_mem.memory_registry.get_op("user-name")
 
@@ -63,7 +63,7 @@ def ops(self) -> list[op_def.OpDef]:
 # objects.
 @weave.op(name="op-ecosystem", render_info={"type": "function"})
 def ecosystem() -> Ecosystem:
-    from weave import registry_mem
+    from weave.legacy import registry_mem
 
     return Ecosystem(
         _orgs=[],
diff --git a/weave/legacy/ecosystem/wandb/wandb_objs.py b/weave/legacy/ecosystem/wandb/wandb_objs.py
index c57d8a296aa..449d82ed786 100644
--- a/weave/legacy/ecosystem/wandb/wandb_objs.py
+++ b/weave/legacy/ecosystem/wandb/wandb_objs.py
@@ -1,7 +1,7 @@
 import typing
 
 import weave
-from weave import registry_mem
+from weave.legacy import registry_mem
 from weave.legacy.ops_domain import run_ops, wb_domain_types
 
 # We can't chain ops called .name() because of a weird bug :( [its a field on VarNode].
diff --git a/weave/legacy/execute.py b/weave/legacy/execute.py
index e05958a2f9c..ccd0d67f101 100644
--- a/weave/legacy/execute.py
+++ b/weave/legacy/execute.py
@@ -15,7 +15,6 @@
     errors,
     parallelism,
     ref_base,
-    registry_mem,
     trace_local,
 )
 from weave import weave_types as types
@@ -27,6 +26,7 @@
 # Trace / cache
 # Language Features
 from weave.legacy import (
+    registry_mem,
     box,
     compile,
     context,
diff --git a/weave/legacy/execute_fast.py b/weave/legacy/execute_fast.py
index 29488de46f2..dbf15feeadd 100644
--- a/weave/legacy/execute_fast.py
+++ b/weave/legacy/execute_fast.py
@@ -4,11 +4,11 @@
     engine_trace,
     errors,
     ref_base,
-    registry_mem,
     weave_internal,
 )
 from weave import weave_types as types
 from weave.legacy import (
+    registry_mem,
     box,
     compile,
     forward_graph,
diff --git a/weave/legacy/language_features/tagging/process_opdef_output_type.py b/weave/legacy/language_features/tagging/process_opdef_output_type.py
index 7364ebcd409..a445e748527 100644
--- a/weave/legacy/language_features/tagging/process_opdef_output_type.py
+++ b/weave/legacy/language_features/tagging/process_opdef_output_type.py
@@ -5,7 +5,7 @@
 
 import typing
 
-from weave import registry_mem
+from weave.legacy import registry_mem
 from weave import weave_types as types
 from weave.legacy import graph
 from weave.legacy.language_features.tagging.opdef_util import (
diff --git a/weave/legacy/op_def_type.py b/weave/legacy/op_def_type.py
index 964aeb4cab9..bb5fccf4525 100644
--- a/weave/legacy/op_def_type.py
+++ b/weave/legacy/op_def_type.py
@@ -11,9 +11,9 @@
 from _ast import AsyncFunctionDef, ExceptHandler
 from typing import Any
 
-from weave import environment, errors, registry_mem, storage
+from weave import environment, errors, storage
 from weave import weave_types as types
-from weave.legacy import artifact_fs, artifact_local, context_state, infer_types
+from weave.legacy import artifact_fs, artifact_local, context_state, infer_types, registry_mem
 
 if typing.TYPE_CHECKING:
     from .op_def import OpDef
diff --git a/weave/legacy/ops_arrow/vectorize.py b/weave/legacy/ops_arrow/vectorize.py
index f4123b18a54..8f3a11a2452 100644
--- a/weave/legacy/ops_arrow/vectorize.py
+++ b/weave/legacy/ops_arrow/vectorize.py
@@ -8,10 +8,9 @@
 
 from weave import (
     errors,
-    registry_mem,
     weave_internal,
 )
-from weave.legacy import weavify
+from weave.legacy import weavify, registry_mem
 from weave import weave_types as types
 from weave.query_api import op, use
 from weave.legacy import dispatch, graph, graph_debug, op_args, op_def
diff --git a/weave/legacy/ops_domain/run_history/history_op_common.py b/weave/legacy/ops_domain/run_history/history_op_common.py
index 1987645a7c9..72d945d39f0 100644
--- a/weave/legacy/ops_domain/run_history/history_op_common.py
+++ b/weave/legacy/ops_domain/run_history/history_op_common.py
@@ -7,12 +7,12 @@
 from weave import (
     engine_trace,
     errors,
-    registry_mem,
     util,
 )
 from weave import weave_types as types
 from weave.query_api import use
 from weave.legacy import (
+    registry_mem,
     _dict_utils,
     artifact_base,
     artifact_fs,
diff --git a/weave/legacy/ops_primitives/weave_api.py b/weave/legacy/ops_primitives/weave_api.py
index 5417cb96de7..33ace6416b4 100644
--- a/weave/legacy/ops_primitives/weave_api.py
+++ b/weave/legacy/ops_primitives/weave_api.py
@@ -5,13 +5,13 @@
 from weave import (
     errors,
     ref_base,
-    registry_mem,
     storage,
     weave_internal,
 )
 from weave import weave_types as types
 from weave.query_api import mutation, op, weave_class
 from weave.legacy import (
+    registry_mem,
     artifact_fs,
     artifact_local,
     artifact_wandb,
diff --git a/weave/legacy/panels_py/generator_templates.py b/weave/legacy/panels_py/generator_templates.py
index 80bcccfb04f..b7fc045cb22 100644
--- a/weave/legacy/panels_py/generator_templates.py
+++ b/weave/legacy/panels_py/generator_templates.py
@@ -18,8 +18,8 @@
 import dataclasses
 import typing
 
-from weave import registry_mem, weave_types
-from weave.legacy import decorator_op, graph
+from weave import weave_types
+from weave.legacy import decorator_op, graph, registry_mem
 
 
 @dataclasses.dataclass
diff --git a/weave/legacy/propagate_gql_keys.py b/weave/legacy/propagate_gql_keys.py
index 50f7d284acc..08aa87a1aff 100644
--- a/weave/legacy/propagate_gql_keys.py
+++ b/weave/legacy/propagate_gql_keys.py
@@ -1,8 +1,8 @@
 import typing
 
-from weave import registry_mem
+
 from weave import weave_types as types
-from weave.legacy import gql_op_plugin, graph, input_provider, op_def, partial_object
+from weave.legacy import gql_op_plugin, graph, input_provider, op_def, partial_object, registry_mem
 
 
 def _propagate_gql_keys_for_node(
diff --git a/weave/registry_mem.py b/weave/legacy/registry_mem.py
similarity index 92%
rename from weave/registry_mem.py
rename to weave/legacy/registry_mem.py
index 29061705502..62f4a42134f 100644
--- a/weave/registry_mem.py
+++ b/weave/legacy/registry_mem.py
@@ -28,7 +28,7 @@ class Registry:
     # the registry over HTTP.
     _updated_at: float
 
-    def __init__(self):
+    def __init__(self):  # type: ignore
         self._types = {}
         self._ops = {}
         self._ops_by_common_name = {}
@@ -41,8 +41,8 @@ def mark_updated(self) -> None:
     def updated_at(self) -> float:
         return self._updated_at
 
-    def register_op(self, op: "OpDef", location=None):
-        if context_state.get_no_op_register():
+    def register_op(self, op: "OpDef", location=None):  # type: ignore
+        if context_state.get_no_op_register():  # type: ignore[no-untyped-call]
             return op
         self.mark_updated()
         # do not save built-in ops today
@@ -97,7 +97,7 @@ def get_op(self, uri: str) -> "OpDef":
             raise errors.WeaveMissingOpDefError("Op not registered: %s" % uri)
         return res
 
-    def find_op_by_fn(self, lazy_local_fn):
+    def find_op_by_fn(self, lazy_local_fn):  # type: ignore
         for op_def in self._op_versions.values():
             if op_def.call_fn == lazy_local_fn:
                 return op_def
@@ -111,7 +111,7 @@ def find_ops_by_common_name(self, common_name: str) -> typing.List["OpDef"]:
         return ops
 
     def find_chainable_ops(self, arg0_type: weave_types.Type) -> typing.List["OpDef"]:
-        def is_chainable(op):
+        def is_chainable(op):  # type: ignore
             if not isinstance(op.input_type, op_args.OpNamedArgs):
                 return False
             args = list(op.input_type.arg_types.values())
@@ -119,9 +119,9 @@ def is_chainable(op):
                 return False
             return args[0].assign_type(arg0_type)
 
-        return [op for op in self._ops.values() if is_chainable(op)]
+        return [op for op in self._ops.values() if is_chainable(op)]  # type: ignore[no-untyped-call]
 
-    def load_saved_ops(self):
+    def load_saved_ops(self):  # type: ignore
         from weave.legacy import op_def_type
 
         for op_ref in storage.objects(op_def_type.OpDefType()):
@@ -149,7 +149,7 @@ def list_packages(self) -> typing.List["OpDef"]:
         ]
         return packages
 
-    def rename_op(self, name, new_name):
+    def rename_op(self, name, new_name):  # type: ignore
         """Internal use only, used during op bootstrapping at decorator time"""
         self.mark_updated()
         op = self._ops.pop(name)
@@ -186,4 +186,4 @@ def rename_op(self, name, new_name):
 
 
 # Processes have a singleton MemoryRegistry
-memory_registry = Registry()
+memory_registry = Registry()  # type: ignore[no-untyped-call]
diff --git a/weave/legacy/stitch.py b/weave/legacy/stitch.py
index 1fd73868f8a..3c2de9a04bc 100644
--- a/weave/legacy/stitch.py
+++ b/weave/legacy/stitch.py
@@ -22,10 +22,10 @@
 import dataclasses
 import typing
 
-from weave.legacy import graph, op_def
+from weave.legacy import graph, op_def, registry_mem
 from weave.legacy.language_features.tagging import opdef_util
 
-from .. import errors, registry_mem
+from .. import errors
 from .. import weave_types as types
 from . import _dict_utils
 
diff --git a/weave/tests/legacy/test_derive_op.py b/weave/tests/legacy/test_derive_op.py
index c3bc56848be..85b86659543 100644
--- a/weave/tests/legacy/test_derive_op.py
+++ b/weave/tests/legacy/test_derive_op.py
@@ -1,5 +1,5 @@
 from ... import api as weave
-from ... import registry_mem
+from ...legacy import registry_mem
 
 
 def test_mapped_add():
diff --git a/weave/tests/legacy/test_mappability.py b/weave/tests/legacy/test_mappability.py
index 9beb8445a63..18822a65478 100644
--- a/weave/tests/legacy/test_mappability.py
+++ b/weave/tests/legacy/test_mappability.py
@@ -2,8 +2,8 @@
 from weave.legacy import context_state as _context
 from weave.legacy import graph
 
-from ... import registry_mem
 from ... import weave_types as types
+from ...legacy import registry_mem
 from ...weave_internal import make_const_node
 
 _loading_builtins_token = _context.set_loading_built_ins()
diff --git a/weave/tests/legacy/test_op_behaviors.py b/weave/tests/legacy/test_op_behaviors.py
index 263f8541647..e49ff51c9c1 100644
--- a/weave/tests/legacy/test_op_behaviors.py
+++ b/weave/tests/legacy/test_op_behaviors.py
@@ -25,7 +25,8 @@
     TaggedValueType,
 )
 
-from ... import registry_mem, storage, weave_internal
+from ... import storage, weave_internal
+from ...legacy import registry_mem
 from ..concrete_tagged_value import (
     TaggedValue,
     concrete_from_tagstore,
diff --git a/weave/tests/legacy/test_op_coverage.py b/weave/tests/legacy/test_op_coverage.py
index 2eadbe3c55d..cede44a1109 100644
--- a/weave/tests/legacy/test_op_coverage.py
+++ b/weave/tests/legacy/test_op_coverage.py
@@ -1,4 +1,4 @@
-from ... import registry_mem
+from ...legacy import registry_mem
 
 
 def make_error_message(missing_ops, section_name):
diff --git a/weave/tests/legacy/test_op_serialization.py b/weave/tests/legacy/test_op_serialization.py
index 66ad75c99d5..7c08a59b982 100644
--- a/weave/tests/legacy/test_op_serialization.py
+++ b/weave/tests/legacy/test_op_serialization.py
@@ -1,8 +1,9 @@
 import pytest
 
 import weave
+import weave.legacy
 
-ops = weave.registry_mem.memory_registry.list_ops()
+ops = weave.legacy.registry_mem.memory_registry.list_ops()
 
 
 def output_type_dict_is_const_function_node(output_type_dict):
diff --git a/weave/tests/legacy/test_serialize.py b/weave/tests/legacy/test_serialize.py
index 85b34ced7f9..3ce4594856f 100644
--- a/weave/tests/legacy/test_serialize.py
+++ b/weave/tests/legacy/test_serialize.py
@@ -2,9 +2,9 @@
 
 import weave
 from weave import query_api as api
-from weave import registry_mem, weave_internal
+from weave import weave_internal
 from weave import weave_types as types
-from weave.legacy import graph, op_args, ops, serialize
+from weave.legacy import graph, op_args, ops, registry_mem, serialize
 from weave.legacy.ops_primitives import list_
 from weave.weave_internal import make_const_node
 
diff --git a/weave/tests/legacy/test_wb_domain_ops.py b/weave/tests/legacy/test_wb_domain_ops.py
index ec2d5bb74bd..5674b65204e 100644
--- a/weave/tests/legacy/test_wb_domain_ops.py
+++ b/weave/tests/legacy/test_wb_domain_ops.py
@@ -9,7 +9,7 @@
 from weave.legacy.ops_domain import wb_domain_types
 from weave.legacy.ops_primitives import _dict_utils
 
-from ... import registry_mem
+from ...legacy import registry_mem
 from .. import fixture_fakewandb as fwb
 
 """
diff --git a/weave/weave_internal.py b/weave/weave_internal.py
index 68a7edbe53c..2b6e07b294b 100644
--- a/weave/weave_internal.py
+++ b/weave/weave_internal.py
@@ -150,7 +150,7 @@ def define_fn(
 # and the function doesn't explicitly operate on tagged values. this ensures that the input tags
 # are propagated appropriately to the output type of the function.
 def refine_graph(node: graph.Node) -> graph.Node:
-    from .registry_mem import memory_registry
+    from .legacy.registry_mem import memory_registry
 
     if isinstance(node, (graph.ConstNode, graph.VoidNode, graph.VarNode)):
         return node
diff --git a/weave/weave_server.py b/weave/weave_server.py
index db96977c0a4..461fe55989d 100644
--- a/weave/weave_server.py
+++ b/weave/weave_server.py
@@ -31,12 +31,18 @@
     errors,
     filesystem,
     logs,
-    registry_mem,
     server,
     storage,
     util,
 )
-from weave.legacy import context_state, graph, value_or_error, wandb_api, weavejs_fixes
+from weave.legacy import (
+    context_state,
+    graph,
+    registry_mem,
+    value_or_error,
+    wandb_api,
+    weavejs_fixes,
+)
 from weave.legacy.language_features.tagging import tag_store
 from weave.server_error_handling import client_safe_http_exceptions_as_werkzeug