Improve typechecking (#357)

droserasprout · web-flow · commit cefdcdd31a51 · 2022-05-26T11:54:52.000Z
diff --git a/src/demo_hic_et_nunc/handlers/on_collect.py b/src/demo_hic_et_nunc/handlers/on_collect.py
@@ -23,7 +23,7 @@ async def on_collect(
     )
     await trade.save()
 
-    swap.amount_left -= int(collect.parameter.objkt_amount)  # type: ignore
+    swap.amount_left -= int(collect.parameter.objkt_amount)
     if swap.amount_left == 0:
         swap.status = models.SwapStatus.FINISHED
     await swap.save()
diff --git a/src/demo_hic_et_nunc/handlers/on_swap.py b/src/demo_hic_et_nunc/handlers/on_swap.py
@@ -11,7 +11,7 @@ async def on_swap(
 ) -> None:
     holder, _ = await models.Holder.get_or_create(address=swap.data.sender_address)
     swap_model = models.Swap(
-        id=int(swap.storage.swap_id) - 1,  # type: ignore
+        id=int(swap.storage.swap_id) - 1,
         creator=holder,
         price=swap.parameter.xtz_per_objkt,
         amount=swap.parameter.objkt_amount,
diff --git a/src/demo_quipuswap/handlers/on_fa12_divest_liquidity.py b/src/demo_quipuswap/handlers/on_fa12_divest_liquidity.py
@@ -40,7 +40,7 @@ async def on_fa12_divest_liquidity(
     share_px = (tez_qty + price * token_qty) / shares_qty
 
     position.realized_pl += shares_qty * (share_px - position.avg_share_px)
-    position.shares_qty -= shares_qty  # type: ignore
+    position.shares_qty -= shares_qty
     assert position.shares_qty >= 0, divest_liquidity.data.hash
 
     await position.save()
diff --git a/src/demo_quipuswap/handlers/on_fa12_invest_liquidity.py b/src/demo_quipuswap/handlers/on_fa12_invest_liquidity.py
@@ -20,19 +20,21 @@ async def on_fa12_invest_liquidity(
     symbol = ctx.template_values['symbol']
     trader = invest_liquidity.data.sender_address
 
+    assert trader is not None
+
     position, _ = await models.Position.get_or_create(trader=trader, symbol=symbol)
 
     assert invest_liquidity.data.amount is not None
     tez_qty = Decimal(invest_liquidity.data.amount) / (10**6)
     token_qty = Decimal(transfer.parameter.value) / (10**decimals)
-    new_shares_qty = int(storage.storage.ledger[trader].balance) + int(storage.storage.ledger[trader].frozen_balance)  # type: ignore
+    new_shares_qty = int(storage.storage.ledger[trader].balance) + int(storage.storage.ledger[trader].frozen_balance)
 
     price = (Decimal(storage.storage.tez_pool) / (10**6)) / (Decimal(storage.storage.token_pool) / (10**decimals))
     value = tez_qty + price * token_qty
     share_px = value / (new_shares_qty - position.shares_qty)
     assert share_px > 0, invest_liquidity.data.hash
 
     position.avg_share_px = (position.shares_qty * position.avg_share_px + value) / new_shares_qty
-    position.shares_qty = new_shares_qty  # type: ignore
+    position.shares_qty = new_shares_qty
 
     await position.save()
diff --git a/src/demo_quipuswap/handlers/on_fa12_transfer.py b/src/demo_quipuswap/handlers/on_fa12_transfer.py
@@ -15,11 +15,11 @@ async def on_fa12_transfer(
     value = int(transfer.parameter.value)
 
     from_position, _ = await models.Position.get_or_create(trader=from_address, symbol=symbol)
-    from_position.shares_qty -= value  # type: ignore
+    from_position.shares_qty -= value
     assert from_position.shares_qty >= 0, transfer.data.hash
     await from_position.save()
 
     to_position, _ = await models.Position.get_or_create(trader=to_address, symbol=symbol)
-    to_position.shares_qty += value  # type: ignore
+    to_position.shares_qty += value
     assert to_position.shares_qty >= 0, transfer.data.hash
     await to_position.save()
diff --git a/src/demo_quipuswap/handlers/on_fa12_withdraw_profit.py b/src/demo_quipuswap/handlers/on_fa12_withdraw_profit.py
@@ -20,6 +20,6 @@ async def on_fa12_withdraw_profit(
     position, _ = await models.Position.get_or_create(trader=trader, symbol=symbol)
     if transaction_0:
         assert transaction_0.amount is not None
-        position.realized_pl += Decimal(transaction_0.amount) / (10**6)  # type: ignore
+        position.realized_pl += Decimal(transaction_0.amount) / (10**6)
 
         await position.save()
diff --git a/src/demo_quipuswap/handlers/on_fa2_divest_liquidity.py b/src/demo_quipuswap/handlers/on_fa2_divest_liquidity.py
@@ -40,7 +40,7 @@ async def on_fa2_divest_liquidity(
     share_px = (tez_qty + price * token_qty) / shares_qty
 
     position.realized_pl += shares_qty * (share_px - position.avg_share_px)
-    position.shares_qty -= shares_qty  # type: ignore
+    position.shares_qty -= shares_qty
     assert position.shares_qty >= 0, divest_liquidity.data.hash
 
     await position.save()
diff --git a/src/demo_quipuswap/handlers/on_fa2_invest_liquidity.py b/src/demo_quipuswap/handlers/on_fa2_invest_liquidity.py
@@ -20,19 +20,21 @@ async def on_fa2_invest_liquidity(
     symbol = ctx.template_values['symbol']
     trader = invest_liquidity.data.sender_address
 
+    assert trader is not None
+
     position, _ = await models.Position.get_or_create(trader=trader, symbol=symbol)
 
     assert invest_liquidity.data.amount is not None
     tez_qty = Decimal(invest_liquidity.data.amount) / (10**6)
     token_qty = sum(Decimal(tx.amount) for tx in transfer.parameter.__root__[0].txs) / (10**decimals)
-    new_shares_qty = int(storage.storage.ledger[trader].balance) + int(storage.storage.ledger[trader].frozen_balance)  # type: ignore
+    new_shares_qty = int(storage.storage.ledger[trader].balance) + int(storage.storage.ledger[trader].frozen_balance)
 
     price = (Decimal(storage.storage.tez_pool) / (10**6)) / (Decimal(storage.storage.token_pool) / (10**decimals))
     value = tez_qty + price * token_qty
     share_px = value / (new_shares_qty - position.shares_qty)
     assert share_px > 0, invest_liquidity.data.hash
 
     position.avg_share_px = (position.shares_qty * position.avg_share_px + value) / new_shares_qty
-    position.shares_qty = new_shares_qty  # type: ignore
+    position.shares_qty = new_shares_qty
 
     await position.save()
diff --git a/src/demo_quipuswap/handlers/on_fa2_transfer.py b/src/demo_quipuswap/handlers/on_fa2_transfer.py
@@ -19,11 +19,11 @@ async def on_fa2_transfer(
         to_address = transfer_tx.to_
         value = int(transfer_tx.amount)
 
-        from_position.shares_qty -= value  # type: ignore
+        from_position.shares_qty -= value
         assert from_position.shares_qty >= 0, transfer.data.hash
 
         to_position, _ = await models.Position.get_or_create(trader=to_address, symbol=symbol)
-        to_position.shares_qty += value  # type: ignore
+        to_position.shares_qty += value
         assert to_position.shares_qty >= 0, transfer.data.hash
         await to_position.save()
 
diff --git a/src/demo_quipuswap/handlers/on_fa2_withdraw_profit.py b/src/demo_quipuswap/handlers/on_fa2_withdraw_profit.py
@@ -21,6 +21,6 @@ async def on_fa2_withdraw_profit(
 
     if transaction_0:
         assert transaction_0.amount is not None
-        position.realized_pl += Decimal(transaction_0.amount) / (10**6)  # type: ignore
+        position.realized_pl += Decimal(transaction_0.amount) / (10**6)
 
         await position.save()
diff --git a/src/demo_tezos_domains/handlers/on_storage_diff.py b/src/demo_tezos_domains/handlers/on_storage_diff.py
@@ -4,7 +4,7 @@
 
 
 async def on_storage_diff(ctx: HandlerContext, storage: NameRegistryStorage) -> None:
-    for name, item in storage.store.records.items():  # type: ignore
+    for name, item in storage.store.records.items():
         record_name = bytes.fromhex(name).decode()
         record_path = record_name.split('.')
         ctx.logger.info('Processing `%s`', record_name)
@@ -22,7 +22,7 @@ async def on_storage_diff(ctx: HandlerContext, storage: NameRegistryStorage) ->
                     defaults={
                         'tld_id': record_path[-1],
                         'owner': item.owner,
-                        'expiry': storage.store.expiry_map.get(item.expiry_key) if item.expiry_key else None,  # type: ignore
+                        'expiry': storage.store.expiry_map.get(item.expiry_key) if item.expiry_key else None,
                         'token_id': int(item.tzip12_token_id) if item.tzip12_token_id else None,
                     },
                 )
diff --git a/src/demo_tzbtc/handlers/on_balance_update.py b/src/demo_tzbtc/handlers/on_balance_update.py
@@ -6,8 +6,8 @@
 
 async def on_balance_update(address: str, balance_update: Decimal, timestamp: datetime):
     holder, _ = await models.Holder.get_or_create(address=address)
-    holder.balance += balance_update  # type: ignore
-    holder.turnover += abs(balance_update)  # type: ignore
-    holder.tx_count += 1  # type: ignore
-    holder.last_seen = timestamp  # type: ignore
+    holder.balance += balance_update
+    holder.turnover += abs(balance_update)
+    holder.tx_count += 1
+    holder.last_seen = timestamp
     await holder.save()
diff --git a/src/demo_tzbtc_transfers/handlers/on_balance_update.py b/src/demo_tzbtc_transfers/handlers/on_balance_update.py
@@ -8,10 +8,10 @@
 async def on_balance_update(address: str, balance_update: Decimal, timestamp: datetime):
     try:
         holder, _ = await models.Holder.get_or_create(address=address)
-        holder.balance += balance_update  # type: ignore
-        holder.turnover += abs(balance_update)  # type: ignore
-        holder.tx_count += 1  # type: ignore
-        holder.last_seen = timestamp  # type: ignore
+        holder.balance += balance_update
+        holder.turnover += abs(balance_update)
+        holder.tx_count += 1
+        holder.last_seen = timestamp
         await holder.save()
     except InvalidOperation:
         return
diff --git a/src/demo_tzcolors/handlers/on_bid.py b/src/demo_tzcolors/handlers/on_bid.py
@@ -9,6 +9,8 @@ async def on_bid(
     ctx: HandlerContext,
     bid: Transaction[BidParameter, TzcolorsAuctionStorage],
 ) -> None:
+    assert bid.data.amount is not None
+
     auction = await models.Auction.filter(
         id=bid.parameter.__root__,
     ).get()
@@ -23,5 +25,5 @@ async def on_bid(
     ).save()
 
     auction.bidder = bidder
-    auction.bid_amount += bid.data.amount  # type: ignore
+    auction.bid_amount += bid.data.amount
     await auction.save()
diff --git a/src/demo_tzcolors/handlers/on_withdraw.py b/src/demo_tzcolors/handlers/on_withdraw.py
@@ -15,7 +15,7 @@ async def on_withdraw(
 
     # FIXME: Don't do that, returns None when id=0. Bug in Tortoise?
     # token = await auction.token
-    token = await models.Token.filter(id=auction.token_id).get()  # type: ignore
+    token = await models.Token.filter(id=auction.token_id).get()
 
     token.holder = await auction.bidder
     await token.save()
diff --git a/src/demo_tzcolors/models.py b/src/demo_tzcolors/models.py
@@ -25,6 +25,8 @@ class Token(Model):
     timestamp = fields.DatetimeField()
     holder: ForeignKeyFieldInstance[Address] = fields.ForeignKeyField('models.Address', 'tokens')
 
+    token_id: int
+
     class Meta:
         table = 'tokens'
 
@@ -40,6 +42,8 @@ class Auction(Model):
     level = fields.BigIntField()
     timestamp = fields.DatetimeField()
 
+    token_id: int
+
     class Meta:
         table = 'auctions'
 
diff --git a/src/dipdup/__main__.py b/src/dipdup/__main__.py
@@ -1,4 +1,4 @@
 from dipdup.cli import cli
 
 if __name__ == '__main__':
-    cli(prog_name='dipdup', standalone_mode=False)  # type: ignore
+    cli(prog_name='dipdup', standalone_mode=False)
diff --git a/src/dipdup/config.py b/src/dipdup/config.py
@@ -1414,8 +1414,7 @@ def _resolve_template(self, template_config: IndexTemplateConfig) -> None:
             value_regex = r'<[ ]*' + key + r'[ ]*>'
             raw_template = re.sub(value_regex, value, raw_template)
 
-        with suppress(AttributeError):
-            missing_value = re.search(r'<*>', raw_template).search(0)  # type: ignore
+        if missing_value := re.search(r'<*>', raw_template):
             raise ConfigurationError(f'`{template_config.name}` index config is missing required template value `{missing_value}`')
 
         json_template = json.loads(raw_template)
diff --git a/src/dipdup/datasources/tzkt/datasource.py b/src/dipdup/datasources/tzkt/datasource.py
@@ -30,7 +30,7 @@
 
 from pysignalr.client import SignalRClient
 from pysignalr.exceptions import ConnectionError as WebsocketConnectionError
-from pysignalr.messages import CompletionMessage  # type: ignore
+from pysignalr.messages import CompletionMessage
 
 from dipdup.config import HTTPConfig
 from dipdup.config import ResolvedIndexConfigT
diff --git a/src/dipdup/datasources/tzkt/models.py b/src/dipdup/datasources/tzkt/models.py
@@ -4,13 +4,17 @@
 from itertools import groupby
 from typing import Any
 from typing import Dict
+from typing import Hashable
 from typing import Iterable
 from typing import List
 from typing import Optional
 from typing import Tuple
 from typing import Type
+from typing import TypeVar
 from typing import Union
+from typing import cast
 
+from pydantic import BaseModel
 from pydantic import Extra
 from pydantic.error_wrappers import ValidationError
 from typing_extensions import get_args
@@ -22,12 +26,15 @@
 
 IntrospectionError = (KeyError, IndexError, AttributeError)
 
+T = TypeVar('T', Hashable, Type[BaseModel])
 
-def extract_root_outer_type(storage_type: Type) -> Type:
+
+def extract_root_outer_type(storage_type: Type[BaseModel]) -> T:
     """Extract Pydantic __root__ type"""
     root_field = storage_type.__fields__['__root__']
     if root_field.allow_none:
-        return typing.Optional[root_field.type_]  # type: ignore
+        # NOTE: Optional is a magic _SpecialForm
+        return cast(Type[BaseModel], typing.Optional[root_field.type_])
     else:
         return root_field.outer_type_
 
@@ -42,7 +49,7 @@ def is_array_type(storage_type: Type) -> bool:
     # NOTE: Pydantic model with __root__ field subclassing List
     with suppress(*IntrospectionError):
         root_type = extract_root_outer_type(storage_type)
-        return is_array_type(root_type)  # type: ignore
+        return is_array_type(root_type)
 
     # NOTE: Something else
     return False
@@ -57,7 +64,7 @@ def get_list_elt_type(list_type: Type[Any]) -> Type[Any]:
 
     # NOTE: Pydantic model with __root__ field subclassing List
     root_type = extract_root_outer_type(list_type)
-    return get_list_elt_type(root_type)  # type: ignore
+    return get_list_elt_type(root_type)
 
 
 @lru_cache(None)
@@ -70,7 +77,7 @@ def get_dict_value_type(dict_type: Type[Any], key: Optional[str] = None) -> Type
     # NOTE: Pydantic model with __root__ field subclassing Dict
     with suppress(*IntrospectionError):
         root_type = extract_root_outer_type(dict_type)
-        return get_dict_value_type(root_type, key)  # type: ignore
+        return get_dict_value_type(root_type, key)
 
     if key is None:
         raise KeyError('Field name or alias is required for object introspection')
@@ -81,7 +88,7 @@ def get_dict_value_type(dict_type: Type[Any], key: Optional[str] = None) -> Type
         if key in (field.name, field.alias):
             # NOTE: Pydantic does not preserve outer_type_ for Optional
             if field.allow_none:
-                return typing.Optional[field.type_]  # type: ignore
+                return cast(Type[Any], typing.Optional[field.type_])
             else:
                 return field.outer_type_
 
@@ -97,7 +104,7 @@ def unwrap_union_type(union_type: Type) -> Tuple[bool, Tuple[Type, ...]]:
 
     with suppress(*IntrospectionError):
         root_type = extract_root_outer_type(union_type)
-        return unwrap_union_type(root_type)  # type: ignore
+        return unwrap_union_type(root_type)
 
     return False, ()
 
@@ -135,10 +142,10 @@ def _apply_bigmap_diffs(
         return dict_storage
 
 
-def _process_storage(storage: Any, storage_type: Type[Any], bigmap_diffs: Dict[int, Iterable[Dict[str, Any]]]) -> Any:
+def _process_storage(storage: Any, storage_type: T, bigmap_diffs: Dict[int, Iterable[Dict[str, Any]]]) -> Any:
     """Replace bigmap pointers with actual data from diffs"""
     # Check if Union or Optional (== Union[Any, NoneType])
-    is_union, arg_types = unwrap_union_type(storage_type)  # type: ignore
+    is_union, arg_types = unwrap_union_type(storage_type)
     if is_union:
         # NOTE: We have no way but trying every possible branch until first success
         for arg_type in arg_types:
@@ -147,12 +154,12 @@ def _process_storage(storage: Any, storage_type: Type[Any], bigmap_diffs: Dict[i
 
     # NOTE: Bigmap pointer, apply diffs
     if isinstance(storage, int) and type(storage) != storage_type:
-        is_array = is_array_type(storage_type)  # type: ignore
+        is_array = is_array_type(storage_type)
         storage = _apply_bigmap_diffs(storage, bigmap_diffs, is_array)
 
     # NOTE: List, process recursively
     elif isinstance(storage, list):
-        elt_type = get_list_elt_type(storage_type)  # type: ignore
+        elt_type = get_list_elt_type(storage_type)
         for i, _ in enumerate(storage):
             storage[i] = _process_storage(storage[i], elt_type, bigmap_diffs)
 
@@ -163,7 +170,7 @@ def _process_storage(storage: Any, storage_type: Type[Any], bigmap_diffs: Dict[i
 
         for key, value in storage.items():
             try:
-                value_type = get_dict_value_type(storage_type, key)  # type: ignore
+                value_type = get_dict_value_type(storage_type, key)
                 storage[key] = _process_storage(value, value_type, bigmap_diffs)
             except IntrospectionError:
                 if not ignore:
diff --git a/src/dipdup/dipdup.py b/src/dipdup/dipdup.py
diff --git a/src/dipdup/hasura.py b/src/dipdup/hasura.py
diff --git a/src/dipdup/http.py b/src/dipdup/http.py
diff --git a/src/dipdup/models.py b/src/dipdup/models.py
diff --git a/src/dipdup/scheduler.py b/src/dipdup/scheduler.py
diff --git a/src/dipdup/utils/database.py b/src/dipdup/utils/database.py
diff --git a/tests/integration_tests/test_demos.py b/tests/integration_tests/test_demos.py

Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@ async def on_collect(`
`23`	`23`	`)`
`24`	`24`	`await trade.save()`
`25`	`25`
`26`		`- swap.amount_left -= int(collect.parameter.objkt_amount) # type: ignore`
	`26`	`+ swap.amount_left -= int(collect.parameter.objkt_amount)`
`27`	`27`	`if swap.amount_left == 0:`
`28`	`28`	`swap.status = models.SwapStatus.FINISHED`
`29`	`29`	`await swap.save()`