chore(weave): Legacy Refactor pt7 (#2225)

* move to legacy: types_numpy

* move to legacy: registry_mem

* move to legacy: weave_types

weave/legacy/artifact_fs.py

@@ -266,7 +266,7 @@ def type(self) -> types.Type:
 
         ot = self._outer_type
         if self.extra is not None:
-            from weave import types_numpy
+            from weave.legacy import types_numpy
 
             if not types.is_list_like(ot) and isinstance(
                 ot, types_numpy.NumpyArrayType

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/types_numpy.py → weave/legacy/types_numpy.py

@@ -1,29 +1,29 @@
 import numpy as np
 
-from . import weave_types as types
+from .. import weave_types as types
 
 
 # TODO: this doesn't match how extra works for list types...
 class NumpyArrayType(types.Type):
     instance_classes = np.ndarray
     name = "WeaveNDArray"
 
-    def __init__(self, dtype="x", shape=(0,)):
+    def __init__(self, dtype="x", shape=(0,)):  # type: ignore
         self.dtype = dtype
         self.shape = shape
 
     @classmethod
-    def type_of_instance(cls, obj):
+    def type_of_instance(cls, obj):  # type: ignore
         return cls(obj.dtype, obj.shape)
 
-    def _to_dict(self):
+    def _to_dict(self):  # type: ignore
         return {"dtype": str(self.dtype), "shape": self.shape}
 
     @classmethod
-    def from_dict(cls, d):
+    def from_dict(cls, d):  # type: ignore
         return cls(np.dtype(d.get("dtype", "object")), d["shape"])
 
-    def _assign_type_inner(self, next_type):
+    def _assign_type_inner(self, next_type):  # type: ignore
         if not isinstance(next_type, NumpyArrayType):
             return False
         if (
@@ -33,13 +33,13 @@ def _assign_type_inner(self, next_type):
             return False
         return True
 
-    def save_instance(self, obj, artifact, name):
+    def save_instance(self, obj, artifact, name):  # type: ignore
         with artifact.new_file(f"{name}.npz", binary=True) as f:
             np.savez_compressed(f, arr=obj)
 
-    def load_instance(self, artifact, name, extra=None):
+    def load_instance(self, artifact, name, extra=None):  # type: ignore
         with artifact.open(f"{name}.npz", binary=True) as f:
             return np.load(f)["arr"]
 
-    def __str__(self):
+    def __str__(self):  # type: ignore
         return "<NumpyArrayType %s %s>" % (self.dtype, self.shape)

weave/weave_pydantic.py → weave/legacy/weave_pydantic.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel, create_model
 
-from . import weave_types as types
-from .legacy import infer_types
+from .. import weave_types as types
+from . import infer_types
 
 
 def weave_type_to_pydantic(

weave/query_api.py

@@ -24,7 +24,7 @@
 from . import weave_types as types
 
 # needed to enable automatic numpy serialization
-from . import types_numpy as _types_numpy
+from .legacy import types_numpy as _types_numpy
 
 from . import errors
 from weave.legacy.decorators import weave_class, mutation, type

weave/serve_fastapi.py

@@ -17,7 +17,8 @@
 from weave.trace.op import Op
 from weave.trace.refs import ObjectRef
 
-from . import errors, weave_pydantic
+from . import errors
+from .legacy import weave_pydantic
 
 key_cache: cache.LruTimeWindowCache[str, typing.Optional[bool]] = (
     cache.LruTimeWindowCache(datetime.timedelta(minutes=5))