diff --git a/api/services/variable_truncator.py b/api/services/variable_truncator.py index 93a032498d..00aa31650c 100644 --- a/api/services/variable_truncator.py +++ b/api/services/variable_truncator.py @@ -1,9 +1,8 @@ from __future__ import annotations import dataclasses -from abc import ABC, abstractmethod from collections.abc import Mapping -from typing import Any, overload, override +from typing import Any, Protocol, overload, override from configs import dify_config from graphon.file import File @@ -66,14 +65,12 @@ class TruncationResult: truncated: bool -class BaseTruncator(ABC): - @abstractmethod - def truncate(self, segment: Segment) -> TruncationResult: - pass +class BaseTruncator(Protocol): + """Protocol for variable truncation strategies.""" - @abstractmethod - def truncate_variable_mapping(self, v: Mapping[str, Any]) -> tuple[Mapping[str, Any], bool]: - pass + def truncate(self, segment: Segment) -> TruncationResult: ... + + def truncate_variable_mapping(self, v: Mapping[str, Any]) -> tuple[Mapping[str, Any], bool]: ... class VariableTruncator(BaseTruncator): diff --git a/api/tests/unit_tests/services/test_variable_truncator_additional.py b/api/tests/unit_tests/services/test_variable_truncator_additional.py index e9427c4ab3..50644d4294 100644 --- a/api/tests/unit_tests/services/test_variable_truncator_additional.py +++ b/api/tests/unit_tests/services/test_variable_truncator_additional.py @@ -1,4 +1,3 @@ -from collections.abc import Mapping from typing import Any import pytest @@ -7,26 +6,7 @@ from graphon.nodes.variable_assigner.common.helpers import UpdatedVariable from graphon.variables.segments import IntegerSegment, ObjectSegment, StringSegment from graphon.variables.types import SegmentType from services import variable_truncator as truncator_module -from services.variable_truncator import BaseTruncator, TruncationResult, VariableTruncator - - -class _AbstractPassthrough(BaseTruncator): - def truncate(self, segment: Any) -> TruncationResult: - return super().truncate(segment) # type: ignore[misc] - - def truncate_variable_mapping(self, v: Mapping[str, Any]) -> tuple[Mapping[str, Any], bool]: - return super().truncate_variable_mapping(v) # type: ignore[misc] - - -class TestBaseTruncatorContract: - def test_base_truncator_methods_should_execute_abstract_placeholders(self) -> None: - passthrough = _AbstractPassthrough() - - truncate_result = passthrough.truncate(StringSegment(value="x")) - mapping_result = passthrough.truncate_variable_mapping({"a": 1}) - - assert truncate_result is None - assert mapping_result is None +from services.variable_truncator import VariableTruncator class TestVariableTruncatorAdditionalBehavior: