Per wallet application

mypy.ini

@@ -78,6 +78,9 @@ disallow_any_expr = False
 disallow_any_explicit = False
 disallow_any_expr = False
 
+[mypy-rp2.per_wallet_tax_engine]
+disallow_any_expr = False
+
 [mypy-rp2.abstract_country]
 disallow_any_explicit = False
 disallow_any_expr = False

src/rp2/abstract_accounting_method.py

@@ -19,7 +19,7 @@
 
 from enum import Enum
 from heapq import heappop, heappush
-from typing import Dict, List, NamedTuple, Optional, Tuple
+from typing import cast, Dict, List, NamedTuple, Optional, Tuple
 
 from rp2.in_transaction import InTransaction
 from rp2.rp2_decimal import ZERO, RP2Decimal
@@ -100,6 +100,10 @@ def set_from_index(self, from_index: int) -> None:
     def set_to_index(self, to_index: int) -> None:
         self.__to_index = to_index  # pylint: disable=unused-private-member
 
+    @property
+    def accounting_method(self) -> "AbstractAccountingMethod":
+        return self._accounting_method
+
     @property
     def from_index(self) -> int:
         return self.__from_index
@@ -132,12 +136,26 @@ def set_partial_amount(self, acquired_lot: InTransaction, amount: RP2Decimal) ->
     def clear_partial_amount(self, acquired_lot: InTransaction) -> None:
         self.set_partial_amount(acquired_lot, ZERO)
 
+    # Reset partial amounts to their original values and from index to zero.
+    def reset_partial_amounts(self, accounting_method: "AbstractAccountingMethod", original_partial_amounts: Dict[InTransaction, RP2Decimal]) -> None:
+        for current_transaction, original_partial_amount in original_partial_amounts.items():
+            self.set_partial_amount(current_transaction, original_partial_amount)
+        self.set_from_index(0)
+
     def __iter__(self) -> AbstractAccountingMethodIterator:
         return self._accounting_method._create_accounting_method_iterator(self)
 
 
 class ChronologicalAcquiredLotCandidates(AbstractAcquiredLotCandidates):
-    pass
+    def __init__(
+        self,
+        accounting_method: "AbstractAccountingMethod",
+        acquired_lot_list: List[InTransaction],
+        acquired_lot_2_partial_amount: Dict[InTransaction, RP2Decimal],
+    ) -> None:
+        super().__init__(accounting_method, acquired_lot_list, acquired_lot_2_partial_amount)
+        if not isinstance(accounting_method, AbstractChronologicalAccountingMethod):
+            raise RP2TypeError(f"Internal error: accounting_method is not of type AbstractChronologicalAccountingMethod, but of type {type(accounting_method)}")
 
 
 class FeatureBasedAcquiredLotCandidates(AbstractAcquiredLotCandidates):
@@ -150,6 +168,8 @@ def __init__(
         acquired_lot_2_partial_amount: Dict[InTransaction, RP2Decimal],
     ) -> None:
         super().__init__(accounting_method, acquired_lot_list, acquired_lot_2_partial_amount)
+        if not isinstance(accounting_method, AbstractFeatureBasedAccountingMethod):
+            raise RP2TypeError(f"Internal error: accounting_method is not of type AbstractFeatureBasedAccountingMethod, but of type {type(accounting_method)}")
         self.__acquired_lot_heap: List[Tuple[AcquiredLotSortKey, InTransaction]] = []
 
     def set_to_index(self, to_index: int) -> None:
@@ -163,6 +183,21 @@ def set_to_index(self, to_index: int) -> None:
     def acquired_lot_heap(self) -> List[Tuple[AcquiredLotSortKey, InTransaction]]:
         return self.__acquired_lot_heap
 
+    # CAUTION:
+    # - acquired_lot must be the last lot chronologically.
+    # - this operation invalidates any outstanding iterator.
+    def add_acquired_lot(self, acquired_lot: InTransaction) -> None:
+        super().add_acquired_lot(acquired_lot)
+        accounting_method = cast(AbstractFeatureBasedAccountingMethod, self.accounting_method)
+        accounting_method.add_selected_lot_to_heap(self.__acquired_lot_heap, acquired_lot)
+
+    def reset_partial_amounts(self, accounting_method: "AbstractAccountingMethod", original_partial_amounts: Dict[InTransaction, RP2Decimal]) -> None:
+        if not isinstance(accounting_method, AbstractFeatureBasedAccountingMethod):
+            raise RP2TypeError(f"Internal error: accounting_method is not of type AbstractFeatureBasedAccountingMethod, but of type {type(accounting_method)}")
+        super().reset_partial_amounts(accounting_method, original_partial_amounts)
+        for current_transaction, original_partial_amount in original_partial_amounts.items():
+            accounting_method.add_selected_lot_to_heap(self.__acquired_lot_heap, current_transaction)
+
 
 class AbstractAccountingMethod:
     def create_lot_candidates(

src/rp2/accounting_engine.py

@@ -204,21 +204,24 @@ def get_acquired_lot_for_taxable_event(
         acquired_lot_and_index: Optional[_AcquiredLotAndIndex] = self.__acquired_lot_avl.find_max_value_less_than(
             self._get_avl_node_key_with_max_disambiguator(taxable_event.timestamp)
         )
-        if acquired_lot_and_index is not None:
-            if acquired_lot_and_index.acquired_lot != self.__acquired_lot_list[acquired_lot_and_index.index]:
-                raise RP2RuntimeError("Internal error: acquired_lot incongruence in accounting logic")
-            method = self._get_accounting_method(taxable_event.timestamp.year)
-            lot_candidates: Optional[AbstractAcquiredLotCandidates] = self.__years_2_lot_candidates.find_max_value_less_than(taxable_event.timestamp.year)
-            # lot_candidates is 1:1 with acquired_lot_and_index, should always be True
-            if lot_candidates:
-                lot_candidates.set_to_index(acquired_lot_and_index.index)
-                acquired_lot_and_amount: Optional[AcquiredLotAndAmount] = method.seek_non_exhausted_acquired_lot(lot_candidates, new_taxable_event_amount)
-                if acquired_lot_and_amount:
-                    return TaxableEventAndAcquiredLot(
-                        taxable_event=taxable_event,
-                        acquired_lot=acquired_lot_and_amount.acquired_lot,
-                        taxable_event_amount=new_taxable_event_amount,
-                        acquired_lot_amount=acquired_lot_and_amount.amount,
-                    )
-
-        raise AcquiredLotsExhaustedException()
+
+        if acquired_lot_and_index is None:
+            raise AcquiredLotsExhaustedException()
+        if acquired_lot_and_index.acquired_lot != self.__acquired_lot_list[acquired_lot_and_index.index]:
+            raise RP2RuntimeError("Internal error: acquired_lot incongruence in accounting logic")
+        method = self._get_accounting_method(taxable_event.timestamp.year)
+        lot_candidates: Optional[AbstractAcquiredLotCandidates] = self.__years_2_lot_candidates.find_max_value_less_than(taxable_event.timestamp.year)
+        # lot_candidates is 1:1 with acquired_lot_and_index, should always be True
+        if not lot_candidates:
+            raise RP2RuntimeError("Internal error: no lot candidates found for year")
+        lot_candidates.set_to_index(acquired_lot_and_index.index)
+        acquired_lot_and_amount: Optional[AcquiredLotAndAmount] = method.seek_non_exhausted_acquired_lot(lot_candidates, new_taxable_event_amount)
+        if not acquired_lot_and_amount:
+            raise AcquiredLotsExhaustedException()
+
+        return TaxableEventAndAcquiredLot(
+            taxable_event=taxable_event,
+            acquired_lot=acquired_lot_and_amount.acquired_lot,
+            taxable_event_amount=new_taxable_event_amount,
+            acquired_lot_amount=acquired_lot_and_amount.amount,
+        )

src/rp2/balance.py

@@ -21,7 +21,7 @@
 
 from rp2.abstract_entry import AbstractEntry
 from rp2.configuration import Configuration
-from rp2.in_transaction import InTransaction
+from rp2.in_transaction import Account, InTransaction
 from rp2.input_data import InputData
 from rp2.intra_transaction import IntraTransaction
 from rp2.logger import LOGGER
@@ -84,12 +84,6 @@ def __repr__(self) -> str:
         return self.to_string(indent=0, repr_format=True)
 
 
-@dataclass(frozen=True, eq=True)
-class Account:
-    exchange: str
-    holder: str
-
-
 class BalanceSet:
     @classmethod
     def type_check(cls, name: str, instance: "BalanceSet") -> "BalanceSet":

src/rp2/gain_loss.py

@@ -200,4 +200,4 @@ def is_long_term_capital_gains(self) -> bool:
             if not self.taxable_event.is_earning():
                 raise RP2RuntimeError("Internal error: acquired lot is None but taxable event is not an earning")
             return False
-        return (self.taxable_event.timestamp - self.acquired_lot.timestamp).days >= self.configuration.country.get_long_term_capital_gain_period()
+        return (self.taxable_event.timestamp - self.acquired_lot.cost_basis_timestamp).days >= self.configuration.country.get_long_term_capital_gain_period()

src/rp2/in_transaction.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from datetime import datetime
+from dataclasses import dataclass
 from typing import Callable, Dict, List, Optional
 
 from rp2.abstract_entry import AbstractEntry
@@ -23,6 +25,12 @@
 from rp2.rp2_error import RP2TypeError, RP2ValueError
 
 
+@dataclass(frozen=True, eq=True)
+class Account:
+    exchange: str
+    holder: str
+
+
 class InTransaction(AbstractTransaction):
     @classmethod
     def type_check(cls, name: str, instance: AbstractEntry) -> "InTransaction":
@@ -49,6 +57,7 @@ def __init__(
         unique_id: Optional[str] = None,
         notes: Optional[str] = None,
         from_lot: Optional["InTransaction"] = None,
+        cost_basis_timestamp: Optional[str] = None,
     ) -> None:
         super().__init__(configuration, timestamp, asset, transaction_type, spot_price, row, unique_id, notes)
 
@@ -63,8 +72,22 @@ def __init__(
             self.__crypto_in = configuration.type_check_positive_decimal("crypto_in", crypto_in, non_zero=True)
         self.__crypto_fee: RP2Decimal = configuration.type_check_positive_decimal("crypto_fee", crypto_fee) if crypto_fee else ZERO
         self.__fiat_fee: RP2Decimal = configuration.type_check_positive_decimal("fiat_fee", fiat_fee) if fiat_fee else ZERO
+
+        # After a transfer, the per-wallet application model is populated with artificial InTransactions that model the "to" part
+        # of the transfer (these artificial InTransactions are not present in the universal application model). These fields are
+        # only populated in the artificial InTransactions of the per-wallet application model field and describes where the funds
+        # come from (due to the transfer) and go to (after the transfer).
         self.__from_lot: Optional[InTransaction] = InTransaction.type_check("from_lot", from_lot) if from_lot is not None else None
-        self.__to_lots: Dict[str, List[InTransaction]] = {}
+        self.__to_lots: Dict[Account, List[InTransaction]] = {}
+
+        # This field is also used only in the artificial InTransactions of the per-wallet application model. If captures all the
+        # upstream InTransactions that the funds came from and it is used for loop detection (it's a map from
+        # wallet -> InTransaction).
+        self.__originates_from: Dict[Account, InTransaction] = {}
+
+        self.__cost_basis_timestamp: Optional[datetime] = (
+            configuration.type_check_timestamp_from_string("cost_basis_timestamp", cost_basis_timestamp) if cost_basis_timestamp else None
+        )
 
         if spot_price == ZERO:
             raise RP2ValueError(f"{self.asset} {type(self).__name__} ({self.timestamp}, id {self.internal_id}): parameter 'spot_price' cannot be 0")
@@ -148,6 +171,7 @@ def to_string(self, indent: int = 0, repr_format: bool = True, extra_data: Optio
             f"fiat_taxable_amount={self.fiat_taxable_amount:.4f}",
             f"from_lot={self.from_lot.internal_id if self.from_lot is not None else ''}",
             f"to_lots={', '.join(to_lots_string_parts)}",
+            f"cost_basis_timestamp={stringify(self.cost_basis_timestamp.strftime('%Y-%m-%d %H:%M:%S.%f %z'))}",
         ]
         if extra_data:
             class_specific_data.extend(extra_data)
@@ -216,14 +240,28 @@ def fiat_balance_change(self) -> RP2Decimal:
     def is_crypto_fee_defined(self) -> bool:
         return self.crypto_fee > ZERO
 
+    # This is only populated in artificial InTransactions of the per-wallet application model and describes where the funds come from
+    # (i.e. the from part of an IntraTransaction that moved these funds).
     @property
     def from_lot(self) -> Optional["InTransaction"]:
         return self.__from_lot
 
+    # This is only populated in artificial InTransactions of the per-wallet application model and describes where the funds go to after
+    # transfers.
     @property
-    def to_lots(self) -> Dict[str, List["InTransaction"]]:
+    def to_lots(self) -> Dict[Account, List["InTransaction"]]:
         return self.__to_lots
 
+    # This is only populated in artificial InTransactions of the per-wallet application model and describes the upstream chain of
+    # transfers that ended up in these funds being moved here.
+    @property
+    def originates_from(self) -> Dict[Account, "InTransaction"]:
+        return self.__originates_from
+
+    @property
+    def cost_basis_timestamp(self) -> datetime:
+        return self.__cost_basis_timestamp if self.__cost_basis_timestamp else self.timestamp
+
     def is_taxable(self) -> bool:
         return self.transaction_type.is_earn_type()
 

src/rp2/plugin/accounting_method/lifo.py

@@ -25,4 +25,4 @@
 # https://ttlc.intuit.com/community/investments-and-rental-properties/discussion/using-lifo-method-for-cryptocurrency-or-even-stock-cost-basis/00/1433542
 class AccountingMethod(AbstractFeatureBasedAccountingMethod):
     def sort_key(self, lot: InTransaction) -> AcquiredLotSortKey:
-        return AcquiredLotSortKey(ZERO, -lot.timestamp.timestamp(), -lot.row)
+        return AcquiredLotSortKey(ZERO, -lot.cost_basis_timestamp.timestamp(), -lot.row)