move to legacy: registry_mem

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
 

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,

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:

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

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

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,

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
 

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=[],

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].

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,

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,

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 (

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

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

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,

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,

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

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(

weave/registry_mem.py → 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,17 +111,17 @@ 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())
             if not args:
                 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]

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
 

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():

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()

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,

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):

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):

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
 

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
 
 """

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