Improve type annotation (#1029)

2024-08-11 02:44:59 -07:00
parent fcc0f5ed99
commit 9dae407812
6 changed files with 83 additions and 41 deletions
--- a/python/sglang/srt/mem_cache/base_prefix_cache.py
+++ b/python/sglang/srt/mem_cache/base_prefix_cache.py
@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from typing import Callable


 class BasePrefixCache(ABC):
@@ -25,7 +26,7 @@ class BasePrefixCache(ABC):
        pass

    @abstractmethod
-    def evict(self, num_tokens, evict_callback):
+    def evict(self, num_tokens: int, evict_callback: Callable):
        pass

    @abstractmethod
@@ -41,7 +42,7 @@ class BasePrefixCache(ABC):
        pass

    def total_size(self):
-        raise NotImplementedError
+        raise NotImplementedError()

    def pretty_print(self):
-        raise NotImplementedError
+        raise NotImplementedError()