Skip to content
Snippets Groups Projects
Unverified Commit 8dfa56ad authored by Jerry Liu's avatar Jerry Liu Committed by GitHub
Browse files

abstract save function to base index class (#39)

parent 3418b11a
No related branches found
No related tags found
No related merge requests found
"""Base data structure classes.""" """Base data structure classes."""
import json
from abc import abstractmethod from abc import abstractmethod
from typing import Any, Generic, List, Optional, TypeVar, cast from typing import Any, Generic, List, Optional, TypeVar, cast
...@@ -92,6 +93,7 @@ class BaseGPTIndex(Generic[IS]): ...@@ -92,6 +93,7 @@ class BaseGPTIndex(Generic[IS]):
def load_from_disk(cls, save_path: str, **kwargs: Any) -> "BaseGPTIndex": def load_from_disk(cls, save_path: str, **kwargs: Any) -> "BaseGPTIndex":
"""Load from disk.""" """Load from disk."""
@abstractmethod
def save_to_disk(self, save_path: str) -> None: def save_to_disk(self, save_path: str) -> None:
"""Safe to file.""" """Safe to file."""
with open(save_path, "w") as f:
json.dump(self.index_struct.to_dict(), f)
"""Init file."""
"""Init file."""
...@@ -72,8 +72,3 @@ class GPTListIndex(BaseGPTIndex[IndexList]): ...@@ -72,8 +72,3 @@ class GPTListIndex(BaseGPTIndex[IndexList]):
"""Load from disk.""" """Load from disk."""
with open(save_path, "r") as f: with open(save_path, "r") as f:
return cls(index_struct=IndexList.from_dict(json.load(f)), **kwargs) return cls(index_struct=IndexList.from_dict(json.load(f)), **kwargs)
def save_to_disk(self, save_path: str) -> None:
"""Safe to file."""
with open(save_path, "w") as f:
json.dump(self.index_struct.to_dict(), f)
"""Init file."""
...@@ -136,8 +136,3 @@ class GPTTreeIndex(BaseGPTIndex[IndexGraph]): ...@@ -136,8 +136,3 @@ class GPTTreeIndex(BaseGPTIndex[IndexGraph]):
"""Load from disk.""" """Load from disk."""
with open(save_path, "r") as f: with open(save_path, "r") as f:
return cls(index_struct=IndexGraph.from_dict(json.load(f)), **kwargs) return cls(index_struct=IndexGraph.from_dict(json.load(f)), **kwargs)
def save_to_disk(self, save_path: str) -> None:
"""Safe to file."""
with open(save_path, "w") as f:
json.dump(self.index_struct.to_dict(), f)
...@@ -4,4 +4,4 @@ profile = "black" ...@@ -4,4 +4,4 @@ profile = "black"
[tool.mypy] [tool.mypy]
ignore_missing_imports = "True" ignore_missing_imports = "True"
disallow_untyped_defs = "True" disallow_untyped_defs = "True"
exclude = ["notebooks"] exclude = ["notebooks", "build"]
\ No newline at end of file \ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment