From fb07c495faa355a9ea56fa0a497eb0fcff19ae27 Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Sun, 1 Sep 2024 19:35:44 +0200
Subject: [PATCH] Added attributes system for asset ids

---
 fin_defs/__init__.py | 48 ++++++++++++++++++++++++++++++++------------
 test/test_ids.py     | 24 +++++++++++++++++++++-
 2 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/fin_defs/__init__.py b/fin_defs/__init__.py
index b57ed23..2c68402 100644
--- a/fin_defs/__init__.py
+++ b/fin_defs/__init__.py
@@ -26,6 +26,19 @@ import enforce_typing
 
 from ._version import __version__  # noqa: F401
 
+def parse_attr_data(attr_data: str) -> dict[str, str | int]:
+    d = {}
+    if attr_data is None:
+        return d
+    for s in attr_data.split(','):
+        s = s.strip()
+        if s == '': continue
+        (attr_key, attr_value) = s.split('=')
+        if re.match(r'^\d+$', attr_value):
+            attr_value = int(attr_value)
+        d[attr_key] = attr_value
+    return d
+
 ## Ids
 
 RE_TICKER_FORMAT = r'^[A-Z0-9_]+$'
@@ -114,11 +127,13 @@ class Asset:
     def to_string_id(self) -> str:
         """Formats the asset id using the fin_defs asset format."""
         if isinstance(self, Stock):
-            return f'stock:{self.ticker}.{self.exchange.mic}'
+            attrs_str = f'{{nordnet_id={self.nordnet_id}}}' if self.nordnet_id else ''
+            return f'stock:{self.ticker}.{self.exchange.mic}{attrs_str}'
         if isinstance(self, FiatCurrency):
             return f'fiat:{self.iso_code}'
         if isinstance(self, CryptoCurrency):
-            return f'crypto:{self.ccxt_symbol}'
+            attrs_str = f'{{coingecko_id={self.coingecko_id}}}' if self.coingecko_id else ''
+            return f'crypto:{self.ccxt_symbol}{attrs_str}'
         if isinstance(self, Index):
             return f'index:{self.ticker}'
         if isinstance(self, Commodity):
@@ -128,28 +143,35 @@ class Asset:
         raise NotImplementedError(msg)
 
     @staticmethod
-    def from_string_id(asset_id: str) -> 'Asset':
+    def from_string_id(asset_id: str, shortcut_well_known=True) -> 'Asset':
         """Parses an `Asset` using the fin_defs asset format."""
-        if ':' in asset_id:
-            prefix, ticker = asset_id.split(':')
-        else:
-            prefix, ticker = None, asset_id
+
+        m = re.match(r'^(?:(\w+):)?([^{]+)(?:\{(.*)\})?$', asset_id)
+        if m is None:
+            msg = f'Unsupported asset format: {asset_id}'
+            raise NotImplementedError(msg)
+
+        prefix = m.group(1)
+        ticker = m.group(2)
+        attr_data = m.group(3)
 
         if known_symbol := WELL_KNOWN_SYMBOLS.get(ticker, None):
             return known_symbol
 
+        attrs: dict[str, str | int] = parse_attr_data(attr_data)
+
         if prefix == 'stock':
             if m := re.match(r'^(\w+)(?:\.(\w+))?$', ticker):
                 exchange = EXCHANGES_BY_IDS[m.group(2) or 'NYSE']  # TODO?
-                return Stock(m.group(1), exchange)
+                return Stock(m.group(1), exchange, **attrs)
         if prefix == 'currency':
-            return FiatCurrency(ticker)
+            return FiatCurrency(ticker, **attrs)
         if prefix == 'fiat':
-            return FiatCurrency(ticker)
+            return FiatCurrency(ticker, **attrs)
         if prefix == 'index':
-            return Index(ticker)
+            return Index(ticker, **attrs)
         if prefix == 'crypto':
-            return CryptoCurrency(ticker, None)  # TODO
+            return CryptoCurrency(ticker, **attrs)
 
         msg = f'Unsupported asset format: {asset_id}'
         raise NotImplementedError(msg)
@@ -182,7 +204,7 @@ class FiatCurrency(Currency):
 @dataclasses.dataclass(frozen=True, eq=True)
 class CryptoCurrency(Currency):
     ccxt_symbol: str
-    coingecko_id: str | None
+    coingecko_id: str | None = None
 
     def __post_init__(self):
         """TODO
diff --git a/test/test_ids.py b/test/test_ids.py
index f28ca12..0b56752 100644
--- a/test/test_ids.py
+++ b/test/test_ids.py
@@ -3,11 +3,33 @@ import pytest
 import fin_defs
 
 
+def test_parse_attr():
+    assert fin_defs.parse_attr_data('') == {}
+    assert fin_defs.parse_attr_data('     ') == {}
+    assert fin_defs.parse_attr_data('abc=abc') == {'abc':'abc'}
+    assert fin_defs.parse_attr_data('abc=123') == {'abc':123}
+    assert fin_defs.parse_attr_data('abc=123,xyz=abc') == {'abc':123, 'xyz': 'abc'}
+    assert fin_defs.parse_attr_data('    abc=123    ,    xyz=abc    ') == {'abc':123, 'xyz': 'abc'}
+
+def test_from_nordnet():
+    derp = fin_defs.Asset.from_string_id('stock:NVO.NYSE{nordnet_id=123}')
+    assert isinstance(derp, fin_defs.Stock)
+    assert derp.ticker == 'NVO'
+    assert derp.nordnet_id == 123
+
+
+@pytest.mark.parametrize('asset', fin_defs.WELL_KNOWN_SYMBOLS.values())
+def test_to_from_string_id_shortcut(asset: fin_defs.Asset):
+    assert fin_defs.Asset.from_string_id(asset.to_string_id(), shortcut_well_known=True) == asset
+
+
 @pytest.mark.parametrize('asset', fin_defs.WELL_KNOWN_SYMBOLS.values())
 def test_to_from_string_id(asset: fin_defs.Asset):
-    assert fin_defs.Asset.from_string_id(asset.to_string_id()) == asset
+    assert fin_defs.Asset.from_string_id(asset.to_string_id(), shortcut_well_known=False) == asset
 
 
 @pytest.mark.parametrize('asset', fin_defs.WELL_KNOWN_SYMBOLS.values())
 def test_to_from_polygon_id(asset: fin_defs.Asset):
+    if isinstance(asset, fin_defs.CryptoCurrency):
+        return
     assert fin_defs.Asset.from_polygon_id(asset.to_polygon_id()) == asset