From 42574f98a46941614d8a113417eae089e5f76a89 Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Thu, 16 May 2024 00:47:42 +0200
Subject: [PATCH] Stricter types

---
 personal_data/data.py                         |  4 ++--
 personal_data/fetchers/partisia_blockchain.py | 11 ++++++++---
 personal_data/main.py                         |  2 +-
 test/test_parsing.py                          |  2 +-
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/personal_data/data.py b/personal_data/data.py
index 59ab28d..16cc9ca 100644
--- a/personal_data/data.py
+++ b/personal_data/data.py
@@ -1,6 +1,6 @@
 import abc
 import dataclasses
-from collections.abc import Iterator
+from collections.abc import Iterator, Mapping
 from enum import Enum
 
 import requests
@@ -38,5 +38,5 @@ class Scraper(abc.ABC):
         return False
 
     @abc.abstractmethod
-    def scrape(self) -> Iterator[object]:
+    def scrape(self) -> Iterator[Mapping[str, object]]:
         pass
diff --git a/personal_data/fetchers/partisia_blockchain.py b/personal_data/fetchers/partisia_blockchain.py
index 0124919..5b6bf77 100644
--- a/personal_data/fetchers/partisia_blockchain.py
+++ b/personal_data/fetchers/partisia_blockchain.py
@@ -3,8 +3,11 @@ import datetime
 import email.utils
 import json
 import logging
+from collections.abc import Iterator, Mapping
 from decimal import Decimal
 
+from frozendict import frozendict
+
 from personal_data.data import DeduplicateMode, Scraper
 
 from .. import secrets
@@ -61,8 +64,7 @@ class MpcBalance(Scraper):
         json_data, date = self.get_json(url, data=data)
         return json_data['coins']['coins']
 
-    def scrape(self):
-        address = secrets.PBC_ACCOUNT_ADDRESS
+    def scrape_latest_values_for(self, address: str) -> frozendict[str, object]:
         coins = self.determine_coins()
 
         url = URL_ACCOUNT_PLUGIN.format(
@@ -93,4 +95,7 @@ class MpcBalance(Scraper):
             data_point['balance.' + coin_data['symbol']] = native_balance
             del coin_idx, coin_data
 
-        yield data_point
+        return frozendict(data_point)
+
+    def scrape(self) -> Iterator[Mapping[str, object]]:
+        yield self.scrape_latest_values_for(secrets.PBC_ACCOUNT_ADDRESS)
diff --git a/personal_data/main.py b/personal_data/main.py
index 3b5ba52..418b83e 100644
--- a/personal_data/main.py
+++ b/personal_data/main.py
@@ -1,8 +1,8 @@
 import csv
 import datetime
+import decimal
 import io
 import logging
-import decimal
 from decimal import Decimal
 
 import requests
diff --git a/test/test_parsing.py b/test/test_parsing.py
index 9996199..b082da9 100644
--- a/test/test_parsing.py
+++ b/test/test_parsing.py
@@ -1,6 +1,6 @@
 import datetime
-
 from decimal import Decimal
+
 import pytest
 
 from personal_data.main import to_value