From 79b9edc5d611c35d03f422408ed3f441068560c0 Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Thu, 29 May 2025 18:42:13 +0200
Subject: [PATCH] Improved price parsing

---
 fin_defs/data.py        | 41 +++++++++++++++++++++++++++--------------
 fin_defs/parse_price.py | 17 +++++++++++------
 2 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/fin_defs/data.py b/fin_defs/data.py
index bf2be56..4200953 100644
--- a/fin_defs/data.py
+++ b/fin_defs/data.py
@@ -214,16 +214,17 @@ class FiatCurrency(Currency):
 
     @staticmethod
     def from_currency_symbol(symbol: str) -> 'Currency | None':
-        for currency, currency_symbol in CURRENCY_SYMBOLS.items():
-            if currency_symbol == symbol:
+        for currency, currency_display in ASSET_DISPLAY.items():
+            if currency_display.symbol == symbol:
                 return currency
         return None
 
     def to_currency_symbol(self) -> str:
-        return CURRENCY_SYMBOLS[self]
+        return ASSET_DISPLAY[self].symbol
 
 
 FiatCurrency.DKK = FiatCurrency('DKK')
+FiatCurrency.SEK = FiatCurrency('SEK')
 FiatCurrency.NOK = FiatCurrency('NOK')
 FiatCurrency.USD = FiatCurrency('USD')
 FiatCurrency.EUR = FiatCurrency('EUR')
@@ -353,6 +354,7 @@ CURRENCY_CODES = {
         'CHF',
         'ZAR',
         'CZK',
+        'SEK',
     ]
 }
 
@@ -364,15 +366,25 @@ WELL_KNOWN_SYMBOLS = (
     | {'SPX500': SPX, 'SP500': SPX, 'Nasdaq 100': NDX}
 )
 
-CURRENCY_SYMBOLS: dict[Currency, str] = {
-    USD: '$',
-    EUR: '€',
-    FiatCurrency.GBP: '£',
-    BTC: '₿',
-    FiatCurrency.JPY: '¥',
-}
 
-ASSET_PREFIX: dict[Asset, str] = CURRENCY_SYMBOLS  # TODO: Remove at some point.
+@dataclasses.dataclass(frozen=True)
+class AssetDisplay:
+    symbol: str | None
+    prefix: str | None
+    suffix: str | None
+
+UNKNOWN_ASSET_DISPLAY = AssetDisplay(None,None,None)
+
+ASSET_DISPLAY: dict[Asset, AssetDisplay] = {
+    USD: AssetDisplay('$','$', None),
+    EUR: AssetDisplay('€','€', None),
+    FiatCurrency.GBP: AssetDisplay('£','£', None),
+    BTC: AssetDisplay('₿','₿',None),
+    FiatCurrency.JPY: AssetDisplay('¥','¥', None),
+    FiatCurrency.DKK: AssetDisplay(None,None, 'kr.'),
+    FiatCurrency.SEK: AssetDisplay(None,None, 'kr.'),
+    FiatCurrency.NOK: AssetDisplay(None,None, 'kr.'),
+}
 
 EXCHANGES = [
     NYSE,
@@ -501,10 +513,11 @@ class AssetAmount:
         elif abs_amount >= THREE_DECIMALS_UNDER_THIS_AMOUNT:
             specificity = '2'
 
-        prefix = ASSET_PREFIX.get(self.asset, '')
-        return ('{sign}{prefix}{amount:.' + specificity + 'f} {name}').format(
+        display = ASSET_DISPLAY.get(self.asset, UNKNOWN_ASSET_DISPLAY)
+        return ('{sign}{prefix}{amount:.' + specificity + 'f}{suffix} {name}').format(
             sign='-' if (self.amount == self.amount and self.amount < 0) else '',
-            prefix=prefix,
+            prefix=display.prefix or '',
+            suffix=f' {display.suffix}' if display.suffix else '',
             amount=abs_amount,
             name=self.asset.raw_short_name(),
         )
diff --git a/fin_defs/parse_price.py b/fin_defs/parse_price.py
index 3aef97b..eaeee28 100644
--- a/fin_defs/parse_price.py
+++ b/fin_defs/parse_price.py
@@ -4,7 +4,7 @@ from decimal import Decimal
 
 from .data import (
     CURRENCY_CODES,
-    CURRENCY_SYMBOLS,
+    ASSET_DISPLAY,
     Asset,
     AssetAmount,
     FiatCurrency,
@@ -30,7 +30,7 @@ def parse_amount(price: str) -> Decimal:
 RE_CURRENCY_CODES = '(?P<code>' + '|'.join(re.escape(c) for c in CURRENCY_CODES) + ')'
 
 RE_CURRENCY_SYMBOLS = (
-    r'(?P<sym>[' + ''.join(re.escape(c) for c in CURRENCY_SYMBOLS.values()) + '])'
+    r'(?P<sym>[' + ''.join(re.escape(c.symbol) for c in ASSET_DISPLAY.values() if c.symbol) + '])'
 )
 
 RE_SYM_AMOUNT_CODE = re.compile(
@@ -77,6 +77,8 @@ RE_AMOUNT_SUFFIX_KR = re.compile(
 
 VARIANTS_OF_FREE = frozenset(['free', 'gratis', 'gives væk'])
 
+KRONER_USING_ASSETS = {k for k,c in ASSET_DISPLAY.items() if c.suffix == 'kr.'}
+
 
 def parse_price(text: str, default_currency: Asset) -> AssetAmount | None:
     """
@@ -94,6 +96,7 @@ def parse_price(text: str, default_currency: Asset) -> AssetAmount | None:
         return AssetAmount(default_currency, Decimal(0))
 
     code, sym, amount_text = None, None, None
+    currency = None
 
     if m := RE_SYM_AMOUNT_CODE.fullmatch(text):
         code, sym, amount_text = m.group('code'), m.group('sym'), m.group('amount')
@@ -104,13 +107,15 @@ def parse_price(text: str, default_currency: Asset) -> AssetAmount | None:
     elif m := RE_AMOUNT_CODE.fullmatch(text):
         code, amount_text = m.group('code'), m.group('amount')
     elif m := (RE_KR_AMOUNT.fullmatch(text) or RE_AMOUNT_SUFFIX_KR.fullmatch(text)):
-        code, amount_text = 'DKK', m.group('amount')
+        amount_text = m.group('amount')
+        currency = default_currency if default_currency in KRONER_USING_ASSETS else FiatCurrency.DKK
     else:
         logger.debug('Unknown format: %s', text)
         return None
 
-    currency = (
-        CURRENCY_CODES[code.upper()] if code else FiatCurrency.from_currency_symbol(sym)
-    )
+    if currency is None:
+        currency = (
+            CURRENCY_CODES[code.upper()] if code else FiatCurrency.from_currency_symbol(sym)
+        )
     assert currency is not None
     return AssetAmount(currency, parse_amount(amount_text))