Repository: tarohi24/typedflow

Issue:
attributes are not inherited

Discussion:

Patch:
diff --git a/typedflow/counted_cache.py b/typedflow/counted_cache.py
index 2936edb..9b2027b 100644
--- a/typedflow/counted_cache.py
+++ b/typedflow/counted_cache.py
@@ -35,3 +35,6 @@ class CacheTable(Generic[H, T]):
             key: H,
             value: T) -> None:
         self.cache_table[key] = CacheItem(count=self.life, value=value)
+
+    def __len__(self) -> int:
+        return len(self.cache_table)
diff --git a/typedflow/nodes.py b/typedflow/nodes.py
index e5da85c..ff4ba9d 100644
--- a/typedflow/nodes.py
+++ b/typedflow/nodes.py
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 import logging
 from typing import (
     Dict,
@@ -26,11 +26,15 @@ logger = logging.getLogger(__file__)
 TD = TypeVar('TD', bound=TypedDict)
 
 
+@dataclass
 class ProviderNode(Generic[K]):
     """
     LoaderNode or TaskNode
     """
-    def __init__(self):
+    succ_count: int = field(init=False)
+    cache_table: CacheTable[int, Batch[K]] = field(init=False)
+
+    def __post_init__(self):
         self.succ_count: int = 0
         self.cache_table: CacheTable[int, Batch[K]]\
             = CacheTable[int, Batch[K]](life=0)
@@ -44,13 +48,17 @@ class ProviderNode(Generic[K]):
         self.cache_table.life += 1
 
 
+@dataclass
 class ConsumerNode(Generic[T]):
     """
     TaskNode or DumpNode
     Note: this is not defined as a dataclass due to the inheritance problem
     see: https://stackoverflow.com/questions/51575931/class-inheritance-in-python-3-7-dataclasses  # noqa
     """
-    def __init__(self, arg_type: Type[T]):
+    arg_type: Type[T] = field(init=False)  # must be same as T
+    precs: Dict[str, ProviderNode[T]] = field(init=False)
+
+    def __post_init__(self, arg_type: Type[T]):
         self.arg_type: Type[T] = arg_type  # must be same as T
         self.precs: Dict[str, ProviderNode[T]] = dict()
 


Test patch:
diff --git a/typedflow/tests/nodes/test_provider.py b/typedflow/tests/nodes/test_provider.py
new file mode 100644
index 0000000..a48ada6
--- /dev/null
+++ b/typedflow/tests/nodes/test_provider.py
@@ -0,0 +1,20 @@
+from typing import List
+
+import pytest
+
+from typedflow.nodes import LoaderNode
+from typedflow.tasks import DataLoader
+
+
+@pytest.fixture
+def loader_node() -> LoaderNode[str]:
+    lst: List[str] = ['hi', 'hello']
+    loader: DataLoader[str] = DataLoader(gen=lst, batch_size=2)
+    node: LoaderNode[str] = LoaderNode(loader=loader)
+    return node
+
+
+def test_init(loader_node):
+    node = loader_node
+    assert node.succ_count == 0
+    assert len(node.cache_table) == 0