Use pydantic models (#64)

* use pydantic

* fix test cases

* v0.6.0

Author: codereverser

.pre-commit-config.yaml

@@ -1,19 +1,23 @@
-# See https://pre-commit.com for more information
-# See https://pre-commit.com/hooks.html for more hooks
 repos:
--   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.1.0
+  - repo: 'https://github.com/pre-commit/pre-commit-hooks'
+    rev: v4.4.0
     hooks:
-    -   id: trailing-whitespace
-    -   id: end-of-file-fixer
-    -   id: check-yaml
-    -   id: check-added-large-files
--   repo: https://github.com/pycqa/flake8
-    rev: '4.0.1'
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-added-large-files
+  - repo: 'https://github.com/pycqa/flake8'
+    rev: 6.0.0
     hooks:
-    -   id: flake8
--   repo: https://github.com/ambv/black
-    rev: 22.1.0
+      - id: flake8
+  - repo: 'https://github.com/psf/black'
+    rev: 23.1.0
     hooks:
-    - id: black
-      language_version: python3.10
+      - id: black
+        language_version: python3.11
+  - repo: 'https://github.com/pycqa/isort'
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        args: ["--profile", "black", "--filter-files"]
+        name: isort (python)

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## 0.6.0 - 2023-02-20
+- use pydantic models for better data validation
+
 ## 0.5.5 - 2022-08-06
 - bug fix with MuPDF parser
 

casparser/__init__.py

@@ -1,11 +1,11 @@
-from .parsers import read_cas_pdf
-from .analysis import CapitalGainsReport
-from .types import CASParserDataType
 from .__version__ import __version__
+from .analysis import CapitalGainsReport
+from .parsers import read_cas_pdf
+from .types import CASData
 
 __all__ = [
     "read_cas_pdf",
     "__version__",
-    "CASParserDataType",
+    "CASData",
     "CapitalGainsReport",
 ]

casparser/analysis/gains.py

@@ -1,29 +1,30 @@
 import csv
+import io
+import itertools
 from collections import deque
 from dataclasses import dataclass
-from decimal import Decimal
 from datetime import date
-import io
-import itertools
+from decimal import Decimal
 from typing import List, Optional
 
 from dateutil.parser import parse as dateparse
 from dateutil.relativedelta import relativedelta
 
-from casparser.exceptions import IncompleteCASError, GainsError
 from casparser.enums import FundType, GainType, TransactionType
-from casparser.types import CASParserDataType, TransactionDataType
+from casparser.exceptions import GainsError, IncompleteCASError
+from casparser.types import CASData, TransactionData
+
 from .utils import CII, get_fin_year, nav_search
 
 PURCHASE_TXNS = {
-    TransactionType.DIVIDEND_REINVEST.name,
-    TransactionType.PURCHASE.name,
-    TransactionType.PURCHASE_SIP.name,
-    TransactionType.REVERSAL.name,
+    TransactionType.DIVIDEND_REINVEST,
+    TransactionType.PURCHASE,
+    TransactionType.PURCHASE_SIP,
+    TransactionType.REVERSAL,
     # Segregated folios are not supported
     # TransactionType.SEGREGATION.name,
-    TransactionType.SWITCH_IN.name,
-    TransactionType.SWITCH_IN_MERGER.name,
+    TransactionType.SWITCH_IN,
+    TransactionType.SWITCH_IN_MERGER,
 }
 
 SALE_TXNS = {
@@ -87,25 +88,25 @@ class MergedTransaction:
     stt: Decimal = Decimal(0.0)
     tds: Decimal = Decimal(0.0)
 
-    def add(self, txn: TransactionDataType):
-        txn_type = txn["type"]
-        if txn_type in PURCHASE_TXNS and txn["units"] is not None:
-            self.nav = txn["nav"]
-            self.purchase_units += txn["units"]
-            self.purchase += txn["amount"]
-        elif txn_type in SALE_TXNS and txn["units"] is not None:
-            self.nav = txn["nav"]
-            self.sale_units += txn["units"]
-            self.sale += txn["amount"]
-        elif txn_type == TransactionType.STT_TAX.name:
-            self.stt += txn["amount"]
-        elif txn_type == TransactionType.STAMP_DUTY_TAX.name:
-            self.stamp_duty += txn["amount"]
-        elif txn_type == TransactionType.TDS_TAX.name:
-            self.tds += txn["amount"]
-        elif txn_type == TransactionType.SEGREGATION.name:
+    def add(self, txn: TransactionData):
+        txn_type = txn.type
+        if txn_type in PURCHASE_TXNS and txn.units is not None:
+            self.nav = txn.nav
+            self.purchase_units += txn.units
+            self.purchase += txn.amount
+        elif txn_type in SALE_TXNS and txn.units is not None:
+            self.nav = txn.nav
+            self.sale_units += txn.units
+            self.sale += txn.amount
+        elif txn_type == TransactionType.STT_TAX:
+            self.stt += txn.amount
+        elif txn_type == TransactionType.STAMP_DUTY_TAX:
+            self.stamp_duty += txn.amount
+        elif txn_type == TransactionType.TDS_TAX:
+            self.tds += txn.amount
+        elif txn_type == TransactionType.SEGREGATION:
             self.nav = Decimal(0.0)
-            self.purchase_units += txn["units"]
+            self.purchase_units += txn.units
             self.purchase = Decimal(0.0)
 
 
@@ -186,7 +187,7 @@ def index_ratio(self) -> Decimal:
 
     @property
     def coa(self) -> Decimal:
-        if self.fund.type == FundType.DEBT.name:
+        if self.fund.type == FundType.DEBT:
             return Decimal(round(self.purchase_value * self.index_ratio, 2))
         if self.purchase_date < self.__cutoff_date:
             if self.sale_date < self.__sell_cutoff_date:
@@ -213,7 +214,7 @@ def stcg(self) -> Decimal:
         return Decimal(0.0)
 
 
-def get_fund_type(transactions: List[TransactionDataType]) -> FundType:
+def get_fund_type(transactions: List[TransactionData]) -> FundType:
     """
     Detect Fund Type.
     - UNKNOWN if there are no redemption transactions
@@ -225,23 +226,23 @@ def get_fund_type(transactions: List[TransactionDataType]) -> FundType:
     """
     valid = any(
         [
-            x["units"] is not None and x["units"] < 0 and x["type"] != TransactionType.REVERSAL.name
+            x.units is not None and x.units < 0 and x.type != TransactionType.REVERSAL
             for x in transactions
         ]
     )
     if not valid:
         return FundType.UNKNOWN
     return (
         FundType.EQUITY
-        if any([x["type"] == TransactionType.STT_TAX.name for x in transactions])
+        if any([x.type == TransactionType.STT_TAX for x in transactions])
         else FundType.DEBT
     )
 
 
 class FIFOUnits:
     """First-In First-Out units calculator."""
 
-    def __init__(self, fund: Fund, transactions: List[TransactionDataType]):
+    def __init__(self, fund: Fund, transactions: List[TransactionData]):
         """
         :param fund: name of fund, mainly for reporting purposes.
         :param transactions: list of transactions for the fund
@@ -264,13 +265,13 @@ def __init__(self, fund: Fund, transactions: List[TransactionDataType]):
     @property
     def clean_transactions(self):
         """remove redundant transactions, without amount"""
-        return filter(lambda x: x["amount"] is not None, self._original_transactions)
+        return filter(lambda x: x.amount is not None, self._original_transactions)
 
     def merge_transactions(self):
         """Group transactions by date with taxes and investments/redemptions separated."""
         merged_transactions = {}
-        for txn in sorted(self.clean_transactions, key=lambda x: (x["date"], -x["amount"])):
-            dt = txn["date"]
+        for txn in sorted(self.clean_transactions, key=lambda x: (x.date, -x.amount)):
+            dt = txn.date
 
             if isinstance(dt, str):
                 dt = dateparse(dt).date()
@@ -347,8 +348,8 @@ def sell(self, sell_date: date, quantity: Decimal, nav: Decimal, tax: Decimal):
 class CapitalGainsReport:
     """Generate Capital Gains Report from the parsed CAS data"""
 
-    def __init__(self, data: CASParserDataType):
-        self._data: CASParserDataType = data
+    def __init__(self, data: CASData):
+        self._data: CASData = data
         self._gains: List[GainEntry] = []
         self.errors = []
         self.invested_amount = Decimal(0.0)
@@ -370,25 +371,25 @@ def get_fy_list(self) -> List[str]:
 
     def process_data(self):
         self._gains = []
-        for folio in self._data.get("folios", []):
-            for scheme in folio.get("schemes", []):
-                transactions = scheme["transactions"]
+        for folio in self._data.folios:
+            for scheme in folio.schemes:
+                transactions = scheme.transactions
                 fund = Fund(
-                    scheme=scheme["scheme"],
-                    folio=folio["folio"],
-                    isin=scheme["isin"],
-                    type=scheme["type"],
+                    scheme=scheme.scheme,
+                    folio=folio.folio,
+                    isin=scheme.isin,
+                    type=scheme.type,
                 )
                 if len(transactions) > 0:
-                    if scheme["open"] >= 0.01:
+                    if scheme.open >= 0.01:
                         raise IncompleteCASError(
                             "Incomplete CAS found. For gains computation, "
                             "all folios should have zero opening balance"
                         )
                     try:
                         fifo = FIFOUnits(fund, transactions)
                         self.invested_amount += fifo.invested
-                        self.current_value += scheme["valuation"]["value"]
+                        self.current_value += scheme.valuation.value
                         self._gains.extend(fifo.gains)
                     except GainsError as exc:
                         self.errors.append((fund.name, str(exc)))

casparser/cli.py

@@ -1,25 +1,25 @@
-from decimal import Decimal
 import itertools
 import os
 import re
 import sys
+from decimal import Decimal
 from typing import Union
 
 import click
 from rich.console import Console
 from rich.markdown import Markdown
 from rich.padding import Padding
-from rich.progress import BarColumn, TextColumn, SpinnerColumn, Progress
-from rich.prompt import Confirm, Prompt
+from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn
+from rich.prompt import Prompt
 from rich.table import Table
 
-from .__version__ import __version__
-
 from . import read_cas_pdf
+from .__version__ import __version__
 from .analysis.gains import CapitalGainsReport
 from .enums import CASFileType
-from .exceptions import ParserException, IncompleteCASError, GainsError
-from .parsers.utils import is_close, cas2json, cas2csv, cas2csv_summary
+from .exceptions import GainsError, IncompleteCASError, ParserException
+from .parsers.utils import cas2csv, cas2csv_summary, cas2json, is_close
+from .types import CASData
 
 CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])
 console = Console()
@@ -38,11 +38,12 @@ def get_color(amount: Union[Decimal, float, int]):
     return "white"
 
 
-def print_summary(data, output_filename=None, include_zero_folios=False):
+def print_summary(parsed_data: CASData, output_filename=None, include_zero_folios=False):
     """Print summary of parsed data."""
     count = 0
     err = 0
 
+    data = parsed_data.dict(by_alias=True)
     is_summary = data["cas_type"] == CASFileType.SUMMARY.name
 
     # Print CAS header stuff
@@ -92,7 +93,6 @@ def print_summary(data, output_filename=None, include_zero_folios=False):
             folio_header_added = False
             current_amc = folio["amc"]
         for scheme in folio["schemes"]:
-
             if scheme["close"] < 1e-3 and not include_zero_folios:
                 continue
 
@@ -130,7 +130,7 @@ def print_summary(data, output_filename=None, include_zero_folios=False):
             count += 1
 
     table = Table(title="Portfolio Summary", show_lines=True)
-    for (hdr, align) in zip(console_header.values(), console_col_align):
+    for hdr, align in zip(console_header.values(), console_col_align):
         # noinspection PyTypeChecker
         table.add_column(hdr, justify=align)
     for row in console_rows:
@@ -152,9 +152,9 @@ def print_summary(data, output_filename=None, include_zero_folios=False):
         console.print(f"File saved : [bold]{output_filename}[/]")
 
 
-def print_gains(data, output_file_path=None, gains_112a=""):
-    cg = CapitalGainsReport(data)
-
+def print_gains(parsed_data: CASData, output_file_path=None, gains_112a=""):
+    cg = CapitalGainsReport(parsed_data)
+    data = parsed_data.dict(by_alias=True)
     if not cg.has_gains():
         console.print("[bold yellow]Warning:[/] No capital gains info found in CAS")
         return
@@ -325,7 +325,7 @@ def cli(output, summary, password, include_all, gains, gains_112a, force_pdfmine
 
     if output_ext in (".csv", ".json"):
         if output_ext == ".csv":
-            if summary or data["cas_type"] == CASFileType.SUMMARY.name:
+            if summary or data.cas_type == CASFileType.SUMMARY.name:
                 description = "Generating summary CSV file..."
                 conv_fn = cas2csv_summary
             else: