refactor

Author: almogdepaz

chia/_tests/plot_sync/test_sync_simulated.py

@@ -69,17 +69,17 @@ async def run(
         initial: bool,
     ) -> None:
         for plot_info in loaded:
-            assert plot_info.prover.get_filepath() not in self.plots
+            assert Path(plot_info.prover.get_filename()) not in self.plots
         for plot_info in removed:
-            assert plot_info.prover.get_filepath() in self.plots
+            assert Path(plot_info.prover.get_filename()) in self.plots
 
         self.invalid = invalid
         self.keys_missing = keys_missing
         self.duplicates = duplicates
 
-        removed_paths: list[Path] = [p.prover.get_filepath() for p in removed] if removed is not None else []
-        invalid_dict: dict[Path, int] = {p.prover.get_filepath(): 0 for p in self.invalid}
-        keys_missing_set: set[Path] = {p.prover.get_filepath() for p in self.keys_missing}
+        removed_paths: list[Path] = [Path(p.prover.get_filename()) for p in removed] if removed is not None else []
+        invalid_dict: dict[Path, int] = {Path(p.prover.get_filename()): 0 for p in self.invalid}
+        keys_missing_set: set[Path] = {Path(p.prover.get_filename()) for p in self.keys_missing}
         duplicates_set: set[str] = {p.prover.get_filename() for p in self.duplicates}
 
         # Inject invalid plots into `PlotManager` of the harvester so that the callback calls below can use them
@@ -91,7 +91,7 @@ async def run(
         # Inject duplicated plots into `PlotManager` of the harvester so that the callback calls below can use them
         # to sync them to the farmer.
         for plot_info in loaded:
-            plot_path = Path(plot_info.prover.get_filepath())
+            plot_path = Path(plot_info.prover.get_filename())
             self.harvester.plot_manager.plot_filename_paths[plot_path.name] = (str(plot_path.parent), set())
         for duplicate in duplicates_set:
             plot_path = Path(duplicate)
@@ -123,9 +123,9 @@ async def sync_done() -> bool:
         await time_out_assert(60, sync_done)
 
         for plot_info in loaded:
-            self.plots[plot_info.prover.get_filepath()] = plot_info
+            self.plots[Path(plot_info.prover.get_filename())] = plot_info
         for plot_info in removed:
-            del self.plots[plot_info.prover.get_filepath()]
+            del self.plots[Path(plot_info.prover.get_filename())]
 
     def validate_plot_sync(self) -> None:
         assert len(self.plots) == len(self.plot_sync_receiver.plots())
@@ -417,16 +417,16 @@ async def test_sync_reset_cases(
         # Inject some data into `PlotManager` of the harvester so that we can validate the reset worked and triggered a
         # fresh sync of all available data of the plot manager
         for plot_info in plots[0:10]:
-            test_data.plots[plot_info.prover.get_filepath()] = plot_info
+            test_data.plots[Path(plot_info.prover.get_filename())] = plot_info
             plot_manager.plots = test_data.plots
         test_data.invalid = plots[10:20]
         test_data.keys_missing = plots[20:30]
         test_data.plot_sync_receiver.simulate_error = simulate_error  # type: ignore[attr-defined]
         sender: Sender = test_runner.test_data[0].plot_sync_sender
         started_sync_id: uint64 = uint64(0)
 
-        plot_manager.failed_to_open_filenames = {p.prover.get_filepath(): 0 for p in test_data.invalid}
-        plot_manager.no_key_filenames = {p.prover.get_filepath() for p in test_data.keys_missing}
+        plot_manager.failed_to_open_filenames = {Path(p.prover.get_filename()): 0 for p in test_data.invalid}
+        plot_manager.no_key_filenames = {Path(p.prover.get_filename()) for p in test_data.keys_missing}
 
         async def wait_for_reset() -> bool:
             assert started_sync_id != 0

chia/_tests/plotting/test_plot_manager.py

@@ -42,9 +42,6 @@
 class MockDiskProver:
     filename: str
 
-    def get_filepath(self) -> Path:
-        return Path(self.filename)
-
     def get_filename(self) -> str:
         return self.filename
 
@@ -115,7 +112,7 @@ def refresh_callback(self, event: PlotRefreshEvents, refresh_result: PlotRefresh
                     for value in actual_value:
                         if type(value) is PlotInfo:
                             for plot_info in expected_list:
-                                if plot_info.prover.get_filename() == value.prover.get_filepath():
+                                if plot_info.prover.get_filename() == value.prover.get_filename():
                                     values_found += 1
                                     continue
                         else:
@@ -507,7 +504,7 @@ async def test_plot_info_caching(environment, bt):
     await refresh_tester.run(expected_result)
     for path, plot_info in env.refresh_tester.plot_manager.plots.items():
         assert path in plot_manager.plots
-        assert plot_manager.plots[path].prover.get_filepath() == plot_info.prover.get_filepath()
+        assert plot_manager.plots[path].prover.get_filename() == plot_info.prover.get_filename()
         assert plot_manager.plots[path].prover.get_id() == plot_info.prover.get_id()
         assert plot_manager.plots[path].prover.get_memo() == plot_info.prover.get_memo()
         assert plot_manager.plots[path].prover.get_size() == plot_info.prover.get_size()
@@ -598,7 +595,7 @@ def assert_cache(expected: list[MockPlotInfo]) -> None:
         test_cache.load()
         assert len(test_cache) == len(expected)
         for plot_info in expected:
-            assert test_cache.get(Path(plot_info.prover.get_filepath())) is not None
+            assert test_cache.get(Path(plot_info.prover.get_filename())) is not None
 
     # Modify two entries, with and without memo modification, they both should remain in the cache after load
     modify_cache_entry(0, 1500, modify_memo=False)
@@ -617,7 +614,7 @@ def assert_cache(expected: list[MockPlotInfo]) -> None:
     # Write the modified cache entries to the file
     cache_path.write_bytes(bytes(VersionedBlob(uint16(CURRENT_VERSION), bytes(cache_data))))
     # And now test that plots in invalid_entries are not longer loaded
-    assert_cache([plot_info for plot_info in plot_infos if str(plot_info.prover.get_filepath()) not in invalid_entries])
+    assert_cache([plot_info for plot_info in plot_infos if plot_info.prover.get_filename() not in invalid_entries])
 
 
 @pytest.mark.anyio

chia/_tests/plotting/test_prover.py

@@ -2,17 +2,17 @@
 
 import tempfile
 from pathlib import Path
+from unittest.mock import MagicMock, patch
 
 import pytest
 
-from chia.plotting.prover import PlotVersion, V2Prover, get_prover_from_file
+from chia.plotting.prover import PlotVersion, V1Prover, V2Prover, get_prover_from_bytes, get_prover_from_file
 
 
 class TestProver:
     def test_v2_prover_init_with_nonexistent_file(self) -> None:
         prover = V2Prover("/nonexistent/path/test.plot2")
         assert prover.get_version() == PlotVersion.V2
-        assert prover.get_filepath() == Path("/nonexistent/path/test.plot2")
         assert prover.get_filename() == "/nonexistent/path/test.plot2"
 
     def test_v2_prover_get_size_raises_error(self) -> None:
@@ -71,6 +71,25 @@ def test_get_prover_from_file_with_plot1_still_works(self) -> None:
             Path(temp_path).unlink()
 
     def test_unsupported_file_extension_raises_value_error(self) -> None:
-        """Test that unsupported file extensions raise ValueError"""
         with pytest.raises(ValueError, match="Unsupported plot file"):
             get_prover_from_file("/nonexistent/path/test.txt")
+
+
+class TestGetProverFromBytes:
+    def test_get_prover_from_bytes_v2_plot(self) -> None:
+        with patch("chia.plotting.prover.V2Prover.from_bytes") as mock_v2_from_bytes:
+            mock_prover = MagicMock()
+            mock_v2_from_bytes.return_value = mock_prover
+            result = get_prover_from_bytes("test.plot2", b"test_data")
+            assert result == mock_prover
+
+    def test_get_prover_from_bytes_v1_plot(self) -> None:
+        with patch("chia.plotting.prover.DiskProver") as mock_disk_prover_class:
+            mock_disk_prover = MagicMock()
+            mock_disk_prover_class.from_bytes.return_value = mock_disk_prover
+            result = get_prover_from_bytes("test.plot", b"test_data")
+            assert isinstance(result, V1Prover)
+
+    def test_get_prover_from_bytes_unsupported_extension(self) -> None:
+        with pytest.raises(ValueError, match="Unsupported plot file"):
+            get_prover_from_bytes("test.txt", b"test_data")

chia/harvester/harvester_api.py

@@ -97,7 +97,7 @@ async def new_signage_point_harvester(
         loop = asyncio.get_running_loop()
 
         def blocking_lookup(filename: Path, plot_info: PlotInfo) -> list[tuple[bytes32, ProofOfSpace]]:
-            # Uses the DiskProver object to lookup qualities. This is a blocking call,
+            # Uses the Prover object to lookup qualities. This is a blocking call,
             # so it should be run in a thread pool.
             try:
                 plot_id = plot_info.prover.get_id()
@@ -218,7 +218,7 @@ def blocking_lookup(filename: Path, plot_info: PlotInfo) -> list[tuple[bytes32,
         async def lookup_challenge(
             filename: Path, plot_info: PlotInfo
         ) -> tuple[Path, list[harvester_protocol.NewProofOfSpace]]:
-            # Executes a DiskProverLookup in a thread pool, and returns responses
+            # Executes a ProverLookup in a thread pool, and returns responses
             all_responses: list[harvester_protocol.NewProofOfSpace] = []
             if self.harvester._shut_down:
                 return filename, []

chia/plotting/cache.py

@@ -55,7 +55,6 @@ class CacheEntry:
 
     @classmethod
     def from_prover(cls, prover: ProverProtocol) -> CacheEntry:
-        """Create CacheEntry from prover"""
         (
             pool_public_key_or_puzzle_hash,
             farmer_public_key,

chia/plotting/manager.py

@@ -386,10 +386,10 @@ def process_file(file_path: Path) -> Optional[PlotInfo]:
                 with self.plot_filename_paths_lock:
                     paths: Optional[tuple[str, set[str]]] = self.plot_filename_paths.get(file_path.name)
                     if paths is None:
-                        paths = (str(Path(cache_entry.prover.get_filepath()).parent), set())
+                        paths = (str(Path(cache_entry.prover.get_filename()).parent), set())
                         self.plot_filename_paths[file_path.name] = paths
                     else:
-                        paths[1].add(str(Path(cache_entry.prover.get_filepath()).parent))
+                        paths[1].add(str(Path(cache_entry.prover.get_filename()).parent))
                         log.warning(f"Have multiple copies of the plot {file_path.name} in {[paths[0], *paths[1]]}.")
                         return None
 
@@ -423,7 +423,7 @@ def process_file(file_path: Path) -> Optional[PlotInfo]:
             plots_refreshed: dict[Path, PlotInfo] = {}
             for new_plot in executor.map(process_file, plot_paths):
                 if new_plot is not None:
-                    plots_refreshed[Path(new_plot.prover.get_filepath())] = new_plot
+                    plots_refreshed[Path(new_plot.prover.get_filename())] = new_plot
             self.plots.update(plots_refreshed)
 
         result.duration = time.time() - start_time

chia/plotting/prover.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
 from enum import IntEnum
-from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar, Protocol, cast
 
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint8
@@ -20,62 +18,30 @@ class PlotVersion(IntEnum):
     V2 = 2
 
 
-class ProverProtocol(ABC):
-    @abstractmethod
-    def get_filepath(self) -> Path:
-        """Returns the filename for the plot"""
-
-    @abstractmethod
-    def get_filename(self) -> str:
-        """Returns the filename string for the plot"""
-
-    @abstractmethod
-    def get_size(self) -> uint8:
-        """Returns the k size of the plot"""
-
-    @abstractmethod
-    def get_memo(self) -> bytes:
-        """Returns the memo"""
-
-    @abstractmethod
-    def get_compression_level(self) -> uint8:
-        """Returns the compression level"""
-
-    @abstractmethod
-    def get_version(self) -> PlotVersion:
-        """Returns the plot version"""
-
-    @abstractmethod
-    def __bytes__(self) -> bytes:
-        """Returns the prover bytes"""
-
-    @abstractmethod
-    def get_id(self) -> bytes32:
-        """Returns the plot ID"""
-
-    @abstractmethod
-    def get_qualities_for_challenge(self, challenge: bytes32) -> list[bytes32]:
-        """Returns the qualities for a given challenge"""
-
-    @abstractmethod
-    def get_full_proof(self, challenge: bytes, index: int, parallel_read: bool = True) -> bytes:
-        """Returns the full proof for a given challenge and index"""
+class ProverProtocol(Protocol):
+    def get_filename(self) -> str: ...
+    def get_size(self) -> uint8: ...
+    def get_memo(self) -> bytes: ...
+    def get_compression_level(self) -> uint8: ...
+    def get_version(self) -> PlotVersion: ...
+    def __bytes__(self) -> bytes: ...
+    def get_id(self) -> bytes32: ...
+    def get_qualities_for_challenge(self, challenge: bytes32) -> list[bytes32]: ...
+    def get_full_proof(self, challenge: bytes, index: int, parallel_read: bool = True) -> bytes: ...
 
     @classmethod
-    @abstractmethod
-    def from_bytes(cls, data: bytes) -> ProverProtocol:
-        """Create a prover from serialized bytes"""
+    def from_bytes(cls, data: bytes) -> ProverProtocol: ...
 
 
-class V2Prover(ProverProtocol):
-    """V2 Plot Prover stubb"""
+class V2Prover:
+    """Placeholder for future V2 plot format support"""
+
+    if TYPE_CHECKING:
+        _protocol_check: ClassVar[ProverProtocol] = cast("V2Prover", None)
 
     def __init__(self, filename: str):
         self._filename = filename
 
-    def get_filepath(self) -> Path:
-        return Path(self._filename)
-
     def get_filename(self) -> str:
         return str(self._filename)
 
@@ -102,29 +68,29 @@ def get_id(self) -> bytes32:
         # TODO: Extract plot ID from V2 plot file
         raise NotImplementedError("V2 plot format is not yet implemented")
 
-    def get_qualities_for_challenge(self, _challenge: bytes) -> list[bytes32]:
+    def get_qualities_for_challenge(self, challenge: bytes) -> list[bytes32]:
         # TODO: todo_v2_plots Implement plot quality lookup
         raise NotImplementedError("V2 plot format is not yet implemented")
 
-    def get_full_proof(self, _challenge: bytes, _index: int, _parallel_read: bool = True) -> bytes:
+    def get_full_proof(self, challenge: bytes, index: int, parallel_read: bool = True) -> bytes:
         # TODO: todo_v2_plots Implement plot proof generation
         raise NotImplementedError("V2 plot format is not yet implemented")
 
     @classmethod
-    def from_bytes(cls, _data: bytes) -> V2Prover:
+    def from_bytes(cls, data: bytes) -> V2Prover:
         # TODO: todo_v2_plots Implement prover deserialization from cache
         raise NotImplementedError("V2 plot format is not yet implemented")
 
 
-class V1Prover(ProverProtocol):
+class V1Prover:
     """Wrapper for existing DiskProver to implement ProverProtocol"""
 
+    if TYPE_CHECKING:
+        _protocol_check: ClassVar[ProverProtocol] = cast("V1Prover", None)
+
     def __init__(self, disk_prover: DiskProver) -> None:
         self._disk_prover = disk_prover
 
-    def get_filepath(self) -> Path:
-        return Path(self._disk_prover.get_filename())
-
     def get_filename(self) -> str:
         return str(self._disk_prover.get_filename())