Add type hints

setup.py

@@ -43,7 +43,7 @@
     install_requires=[
         "asyncio==3.4.3",
         "pyyaml==6.0",
-        "redis==4.4.4",
+        "redis==5.0.3",
         "sgx.py==0.9dev2",
         "skale-contracts==1.0.1a5",
         "typing-extensions==4.9.0",

skale/contracts/allocator/__init__.py

@@ -2,3 +2,5 @@
 
 from skale.contracts.allocator.escrow import Escrow
 from skale.contracts.allocator.allocator import Allocator
+
+__all__ = ['Allocator', 'Escrow']

skale/contracts/allocator/allocator.py

@@ -18,11 +18,24 @@
 #   along with SKALE.py.  If not, see <https://www.gnu.org/licenses/>.
 """ SKALE Allocator Core Escrow methods """
 
-from enum import IntEnum
-
-from skale.contracts.base_contract import BaseContract, transaction_method
-from skale.transactions.exceptions import ContractLogicError
-from skale.transactions.result import TxRes
+from typing import Any, Dict, List
+
+from eth_typing import ChecksumAddress
+from web3 import Web3
+from web3.contract.contract import ContractFunction
+from web3.exceptions import ContractLogicError
+from web3.types import Wei
+
+from skale.contracts.allocator_contract import AllocatorContract
+from skale.contracts.base_contract import transaction_method
+from skale.types.allocation import (
+    BeneficiaryStatus,
+    BeneficiaryPlan,
+    Plan,
+    PlanId,
+    PlanWithId,
+    TimeUnit
+)
 from skale.utils.helper import format_fields
 
 
@@ -47,36 +60,25 @@
 MAX_NUM_OF_BENEFICIARIES = 9999
 
 
-class TimeUnit(IntEnum):
-    DAY = 0
-    MONTH = 1
-    YEAR = 2
-
-
-class BeneficiaryStatus(IntEnum):
-    UNKNOWN = 0
-    CONFIRMED = 1
-    ACTIVE = 2
-    TERMINATED = 3
-
-
-class Allocator(BaseContract):
-    def is_beneficiary_registered(self, beneficiary_address: str) -> bool:
+class Allocator(AllocatorContract):
+    def is_beneficiary_registered(self, beneficiary_address: ChecksumAddress) -> bool:
         """Confirms whether the beneficiary is registered in a Plan.
 
         :returns: Boolean value
         :rtype: bool
         """
-        return self.contract.functions.isBeneficiaryRegistered(beneficiary_address).call()
+        return bool(self.contract.functions.isBeneficiaryRegistered(beneficiary_address).call())
 
-    def is_delegation_allowed(self, beneficiary_address: str) -> bool:
-        return self.contract.functions.isDelegationAllowed(beneficiary_address).call()
+    def is_delegation_allowed(self, beneficiary_address: ChecksumAddress) -> bool:
+        return bool(self.contract.functions.isDelegationAllowed(beneficiary_address).call())
 
-    def is_vesting_active(self, beneficiary_address: str) -> bool:
-        return self.contract.functions.isVestingActive(beneficiary_address).call()
+    def is_vesting_active(self, beneficiary_address: ChecksumAddress) -> bool:
+        return bool(self.contract.functions.isVestingActive(beneficiary_address).call())
 
-    def get_escrow_address(self, beneficiary_address: str) -> str:
-        return self.contract.functions.getEscrowAddress(beneficiary_address).call()
+    def get_escrow_address(self, beneficiary_address: ChecksumAddress) -> ChecksumAddress:
+        return Web3.to_checksum_address(
+            self.contract.functions.getEscrowAddress(beneficiary_address).call()
+        )
 
     @transaction_method
     def add_plan(
@@ -87,7 +89,7 @@
             vesting_interval: int,
             can_delegate: bool,
             is_terminatable: bool
-    ) -> TxRes:
+    ) -> ContractFunction:
         return self.contract.functions.addPlan(
             vestingCliff=vesting_cliff,
             totalVestingDuration=total_vesting_duration,
@@ -100,12 +102,12 @@
     @transaction_method
     def connect_beneficiary_to_plan(
             self,
-            beneficiary_address: str,
+            beneficiary_address: ChecksumAddress,
             plan_id: int,
             start_month: int,
             full_amount: int,
             lockup_amount: int,
-    ) -> TxRes:
+    ) -> ContractFunction:
         return self.contract.functions.connectBeneficiaryToPlan(
             beneficiary=beneficiary_address,
             planId=plan_id,
@@ -115,61 +117,102 @@
         )
 
     @transaction_method
-    def start_vesting(self, beneficiary_address: str) -> TxRes:
+    def start_vesting(self, beneficiary_address: ChecksumAddress) -> ContractFunction:
         return self.contract.functions.startVesting(beneficiary_address)
 
     @transaction_method
-    def stop_vesting(self, beneficiary_address: str) -> TxRes:
+    def stop_vesting(self, beneficiary_address: ChecksumAddress) -> ContractFunction:
         return self.contract.functions.stopVesting(beneficiary_address)
 
     @transaction_method
-    def grant_role(self, role: bytes, address: str) -> TxRes:
+    def grant_role(self, role: bytes, address: ChecksumAddress) -> ContractFunction:
         return self.contract.functions.grantRole(role, address)
 
     def vesting_manager_role(self) -> bytes:
-        return self.contract.functions.VESTING_MANAGER_ROLE().call()
+        return bytes(self.contract.functions.VESTING_MANAGER_ROLE().call())
 
-    def has_role(self, role: bytes, address: str) -> bool:
-        return self.contract.functions.hasRole(role, address).call()
+    def has_role(self, role: bytes, address: ChecksumAddress) -> bool:
+        return bool(self.contract.functions.hasRole(role, address).call())
 
-    def __get_beneficiary_plan_params_raw(self, beneficiary_address: str):
-        return self.contract.functions.getBeneficiaryPlanParams(beneficiary_address).call()
+    def __get_beneficiary_plan_params_raw(self, beneficiary_address: ChecksumAddress) -> List[Any]:
+        return list(self.contract.functions.getBeneficiaryPlanParams(beneficiary_address).call())
 
     @format_fields(BENEFICIARY_FIELDS)
-    def get_beneficiary_plan_params_dict(self, beneficiary_address: str) -> dict:
+    def get_beneficiary_plan_params_dict(self, beneficiary_address: ChecksumAddress) -> List[Any]:
         return self.__get_beneficiary_plan_params_raw(beneficiary_address)
 
-    def get_beneficiary_plan_params(self, beneficiary_address: str) -> dict:
+    def get_beneficiary_plan_params(self, beneficiary_address: ChecksumAddress) -> BeneficiaryPlan:
         plan_params = self.get_beneficiary_plan_params_dict(beneficiary_address)
-        plan_params['statusName'] = BeneficiaryStatus(plan_params['status']).name
-        return plan_params
-
-    def __get_plan_raw(self, plan_id: int):
-        return self.contract.functions.getPlan(plan_id).call()
+        if plan_params is None:
+            raise ValueError('Plan for ', beneficiary_address, ' is missing')
+        if isinstance(plan_params, list):
+            return self._to_beneficiary_plan({
+                **plan_params[0],
+                'statusName': BeneficiaryStatus(plan_params[0]['status']).name
+            })
+        if isinstance(plan_params, dict):
+            return self._to_beneficiary_plan({
+                **plan_params,
+                'statusName': BeneficiaryStatus(plan_params.get('status', 0)).name
+            })
+        raise TypeError(f'Internal error on getting plan params for ${beneficiary_address}')
+
+    def __get_plan_raw(self, plan_id: PlanId) -> List[Any]:
+        return list(self.contract.functions.getPlan(plan_id).call())
 
     @format_fields(PLAN_FIELDS)
-    def get_plan(self, plan_id: int) -> dict:
+    def get_untyped_plan(self, plan_id: PlanId) -> List[Any]:
         return self.__get_plan_raw(plan_id)
 
-    def get_all_plans(self) -> dict:
+    def get_plan(self, plan_id: PlanId) -> Plan:
+        untyped_plan = self.get_untyped_plan(plan_id)
+        if untyped_plan is None:
+            raise ValueError('Plan ', plan_id, ' is missing')
+        if isinstance(untyped_plan, list):
+            return self._to_plan(untyped_plan[0])
+        if isinstance(untyped_plan, dict):
+            return self._to_plan(untyped_plan)
+        raise TypeError(plan_id)
+
+    def get_all_plans(self) -> List[PlanWithId]:
         plans = []
         for i in range(1, MAX_NUM_OF_PLANS):
             try:
-                plan = self.get_plan(i)
-                plan['planId'] = i
+                plan_id = PlanId(i)
+                plan = PlanWithId({**self.get_plan(plan_id), 'planId': plan_id})
                 plans.append(plan)
             except (ContractLogicError, ValueError):
                 break
         return plans
 
-    def calculate_vested_amount(self, address: str) -> int:
-        return self.contract.functions.calculateVestedAmount(address).call()
-
-    def get_finish_vesting_time(self, address: str) -> int:
-        return self.contract.functions.getFinishVestingTime(address).call()
-
-    def get_lockup_period_end_timestamp(self, address: str) -> int:
-        return self.contract.functions.getLockupPeriodEndTimestamp(address).call()
-
-    def get_time_of_next_vest(self, address: str) -> int:
-        return self.contract.functions.getTimeOfNextVest(address).call()
+    def calculate_vested_amount(self, address: ChecksumAddress) -> Wei:
+        return Wei(self.contract.functions.calculateVestedAmount(address).call())
+
+    def get_finish_vesting_time(self, address: ChecksumAddress) -> int:
+        return int(self.contract.functions.getFinishVestingTime(address).call())
+
+    def get_lockup_period_end_timestamp(self, address: ChecksumAddress) -> int:
+        return int(self.contract.functions.getLockupPeriodEndTimestamp(address).call())
+
+    def get_time_of_next_vest(self, address: ChecksumAddress) -> int:
+        return int(self.contract.functions.getTimeOfNextVest(address).call())
+
+    def _to_plan(self, untyped_plan: Dict[str, Any]) -> Plan:
+        return Plan({
+            'totalVestingDuration': int(untyped_plan['totalVestingDuration']),
+            'vestingCliff': int(untyped_plan['vestingCliff']),
+            'vestingIntervalTimeUnit': TimeUnit(untyped_plan['vestingIntervalTimeUnit']),
+            'vestingInterval': int(untyped_plan['vestingInterval']),
+            'isDelegationAllowed': bool(untyped_plan['isDelegationAllowed']),
+            'isTerminatable': bool(untyped_plan['isTerminatable'])
+        })
+
+    def _to_beneficiary_plan(self, untyped_beneficiary_plan: Dict[str, Any]) -> BeneficiaryPlan:
+        return BeneficiaryPlan({
+            'status': BeneficiaryStatus(untyped_beneficiary_plan['status']),
+            'statusName': str(untyped_beneficiary_plan['statusName']),
+            'planId': PlanId(untyped_beneficiary_plan['planId']),
+            'startMonth': int(untyped_beneficiary_plan['startMonth']),
+            'fullAmount': Wei(untyped_beneficiary_plan['fullAmount']),
+            'amountAfterLockup': Wei(untyped_beneficiary_plan['amountAfterLockup'])
+        })

skale/contracts/allocator/escrow.py

@@ -18,32 +18,49 @@
 #   along with SKALE.py.  If not, see <https://www.gnu.org/licenses/>.
 """ SKALE Allocator Core Escrow methods """
 
+from __future__ import annotations
 import functools
+from typing import Any, Callable, TYPE_CHECKING
 
-from skale.contracts.base_contract import BaseContract, transaction_method
+from eth_typing import ChecksumAddress
+from web3.contract.contract import ContractFunction
+from web3.types import Wei
+
+from skale.contracts.allocator_contract import AllocatorContract
+from skale.contracts.base_contract import transaction_method
 from skale.transactions.result import TxRes
+from skale.types.delegation import DelegationId
+from skale.types.validator import ValidatorId
+
+if TYPE_CHECKING:
+    from skale.contracts.allocator.allocator import Allocator
 
 
-def beneficiary_escrow(transaction):
+def beneficiary_escrow(transaction: Callable[..., TxRes]) -> Callable[..., TxRes]:
     @functools.wraps(transaction)
-    def wrapper(self, *args, beneficiary_address, **kwargs):
+    def wrapper(
+            self: AllocatorContract,
+            *args: Any,
+            beneficiary_address: ChecksumAddress,
+            **kwargs: Any
+    ) -> TxRes:
         self.contract = self.skale.instance.get_contract('Escrow', beneficiary_address)
         return transaction(self, *args, **kwargs)
     return wrapper
 
 
-class Escrow(BaseContract):
+class Escrow(AllocatorContract):
     @property
     @functools.lru_cache()
-    def allocator(self):
+    def allocator(self) -> Allocator:
         return self.skale.allocator
 
-    def init_contract(self, skale, address, abi) -> None:
-        self.contract = None
+    def init_contract(self, *args: Any) -> None:
+        self.contract = self.allocator.contract
 
     @beneficiary_escrow
     @transaction_method
-    def retrieve(self) -> TxRes:
+    def retrieve(self) -> ContractFunction:
         """Allows Holder to retrieve vested tokens from the Escrow contract
 
         :returns: Transaction results
@@ -53,7 +70,7 @@
 
     @beneficiary_escrow
     @transaction_method
-    def retrieve_after_termination(self, address: str) -> TxRes:
+    def retrieve_after_termination(self, address: ChecksumAddress) -> ContractFunction:
         """Allows Core Owner to retrieve remaining transferrable escrow balance
         after Core holder termination. Slashed tokens are non-transferable
 
@@ -64,7 +81,13 @@
 
     @beneficiary_escrow
     @transaction_method
-    def delegate(self, validator_id: int, amount: int, delegation_period: int, info: str) -> TxRes:
+    def delegate(
+            self,
+            validator_id: ValidatorId,
+            amount: Wei,
+            delegation_period: int,
+            info: str
+    ) -> ContractFunction:
         """Allows Core holder to propose a delegation to a validator
 
         :param validator_id: ID of the validator to delegate tokens
@@ -82,7 +105,7 @@
 
     @beneficiary_escrow
     @transaction_method
-    def request_undelegation(self, delegation_id: int) -> TxRes:
+    def request_undelegation(self, delegation_id: DelegationId) -> ContractFunction:
         """Allows Holder and Owner to request undelegation. Only Owner can
         request undelegation after Core holder is deactivated (upon holder termination)
 
@@ -95,7 +118,7 @@
 
     @beneficiary_escrow
     @transaction_method
-    def withdraw_bounty(self, validator_id: int, to: str) -> TxRes:
+    def withdraw_bounty(self, validator_id: ValidatorId, to: ChecksumAddress) -> ContractFunction:
         """Allows Beneficiary and Vesting Owner to withdraw earned bounty.
 
         :param validator_id: ID of the validator
@@ -109,7 +132,7 @@
 
     @beneficiary_escrow
     @transaction_method
-    def cancel_pending_delegation(self, delegation_id: int) -> TxRes:
+    def cancel_pending_delegation(self, delegation_id: DelegationId) -> ContractFunction:
         """Cancel pending delegation request.
 
         :param delegation_id: ID of the delegation to cancel

skale/dataclasses/delegation_status.py → skale/contracts/allocator_contract.py

@@ -16,15 +16,9 @@
 #
 #   You should have received a copy of the GNU Affero General Public License
 #   along with SKALE.py.  If not, see <https://www.gnu.org/licenses/>.
+from skale.contracts.base_contract import BaseContract
+from skale.skale_allocator import SkaleAllocator
 
-from enum import Enum
 
-
-class DelegationStatus(Enum):
-    PROPOSED = 0
-    ACCEPTED = 1
-    CANCELED = 2
-    REJECTED = 3
-    DELEGATED = 4
-    UNDELEGATION_REQUESTED = 5
-    COMPLETED = 6
+class AllocatorContract(BaseContract[SkaleAllocator]):
+    pass