From 9a7e0cecf0ed1513f4f97b7bfa66172bc2d66205 Mon Sep 17 00:00:00 2001
From: Aaron <29749331+aarnphm@users.noreply.github.com>
Date: Thu, 30 Nov 2023 09:42:24 -0500
Subject: [PATCH] fix(types): makes sures mypy is running strict

Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
---
 .pre-commit-config.yaml                               | 2 +-
 openllm-core/src/openllm_core/utils/__init__.pyi      | 4 ++++
 openllm-python/src/openllm/__init__.pyi               | 2 --
 openllm-python/src/openllm/bundle/__init__.py         | 3 +--
 openllm-python/src/openllm/bundle/__init__.pyi        | 8 ++++----
 openllm-python/src/openllm/serialisation/_helpers.pyi | 4 ++--
 6 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 430c307e..57032b40 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -22,7 +22,7 @@ repos:
         language: system
         always_run: true
         pass_filenames: false
-        entry: python ./tools/update-mypy.py && mypy --strict
+        entry: mypy --strict
   - repo: https://github.com/editorconfig-checker/editorconfig-checker.python
     rev: '2.7.3'
     hooks:
diff --git a/openllm-core/src/openllm_core/utils/__init__.pyi b/openllm-core/src/openllm_core/utils/__init__.pyi
index 787ee940..c9ab9baa 100644
--- a/openllm-core/src/openllm_core/utils/__init__.pyi
+++ b/openllm-core/src/openllm_core/utils/__init__.pyi
@@ -38,6 +38,10 @@ from openllm_core.utils.serde import converter as converter
 DEBUG: bool = ...
 SHOW_CODEGEN: bool = ...
 MYPY: bool = ...
+DEBUG_ENV_VAR: str = ...
+QUIET_ENV_VAR: str = ...
+DEV_DEBUG_VAR: str = ...
+WARNING_ENV_VAR: str = ...
 
 _T = TypeVar('_T')
 
diff --git a/openllm-python/src/openllm/__init__.pyi b/openllm-python/src/openllm/__init__.pyi
index 3cf1d2fb..3b966a90 100644
--- a/openllm-python/src/openllm/__init__.pyi
+++ b/openllm-python/src/openllm/__init__.pyi
@@ -19,7 +19,6 @@ from openllm_cli._sdk import (
   import_model as import_model,
   list_models as list_models,
   start as start,
-  start_grpc as start_grpc,
 )
 from openllm_core._configuration import (
   GenerationConfig as GenerationConfig,
@@ -37,7 +36,6 @@ from . import (
   client as client,
   exceptions as exceptions,
   serialisation as serialisation,
-  testing as testing,
   utils as utils,
 )
 from ._deprecated import Runner as Runner
diff --git a/openllm-python/src/openllm/bundle/__init__.py b/openllm-python/src/openllm/bundle/__init__.py
index 5dbe6806..aa20f61b 100644
--- a/openllm-python/src/openllm/bundle/__init__.py
+++ b/openllm-python/src/openllm/bundle/__init__.py
@@ -29,8 +29,7 @@ class RefResolver:
   @property
   def tag(self) -> str: return 'latest' if self.strategy in {'latest', 'nightly'} else repr(self.version)
 __lazy = LazyModule(
-  __name__,
-  os.path.abspath('__file__'),
+  __name__, os.path.abspath('__file__'), #
   {'_package': ['create_bento', 'build_editable', 'construct_python_options', 'construct_docker_options']},
   extra_objects={'RefResolver': RefResolver}
 )
diff --git a/openllm-python/src/openllm/bundle/__init__.pyi b/openllm-python/src/openllm/bundle/__init__.pyi
index c09be4c6..024e7650 100644
--- a/openllm-python/src/openllm/bundle/__init__.pyi
+++ b/openllm-python/src/openllm/bundle/__init__.pyi
@@ -2,10 +2,10 @@ from typing import Optional
 
 import attr
 
-from openllm_core._typing_compat import LiteralContainerVersionStrategy
+from openllm_core._typing_compat import LiteralVersionStrategy
 from openllm_core.utils.lazy import VersionInfo
 
-from . import _package as _package, oci as oci
+from . import _package as _package
 from ._package import (
   build_editable as build_editable,
   construct_docker_options as construct_docker_options,
@@ -17,9 +17,9 @@ from ._package import (
 class RefResolver:
   git_hash: str
   version: VersionInfo
-  strategy: LiteralContainerVersionStrategy
+  strategy: LiteralVersionStrategy
 
   @classmethod
-  def from_strategy(cls, strategy_or_version: Optional[LiteralContainerVersionStrategy] = ...) -> RefResolver: ...
+  def from_strategy(cls, strategy_or_version: Optional[LiteralVersionStrategy] = ...) -> RefResolver: ...
   @property
   def tag(self) -> str: ...
diff --git a/openllm-python/src/openllm/serialisation/_helpers.pyi b/openllm-python/src/openllm/serialisation/_helpers.pyi
index a4c5865c..7f2628a1 100644
--- a/openllm-python/src/openllm/serialisation/_helpers.pyi
+++ b/openllm-python/src/openllm/serialisation/_helpers.pyi
@@ -1,6 +1,6 @@
 import types
 from contextlib import contextmanager
-from typing import Optional, Sequence
+from typing import Iterator, Optional, Sequence, Tuple
 
 import transformers
 
@@ -21,4 +21,4 @@ def save_model(
   trust_remote_code: bool,
   module: str,
   external_modules: Sequence[types.ModuleType],
-) -> Model: ...
+) -> Iterator[Tuple[Model, Sequence[types.ModuleType]]]: ...