Refactor: support etf type (#18)

* feat: add etf type in stock proto

* feat: add etf type in domain

* feat: add stock proto in adapter

* feat: add stock proto in usecase

* feat: add stock proto in handler

---------

Co-authored-by: Eyo Chen <eyo.chen@amazingtalker.com>

Author: eyo-chen

src/adapters/portfolio.py

@@ -1,9 +1,9 @@
-from dataclasses import asdict
 from datetime import datetime, timezone
 from pymongo import MongoClient
 from pymongo.database import Database
 from .base import AbstractPortfolioRepository
 from domain.portfolio import Portfolio, Holding
+from domain.enum import StockType
 
 
 class PortfolioRepository(AbstractPortfolioRepository):
@@ -22,7 +22,12 @@ def get(self, user_id: int) -> Portfolio:
             cash_balance=result["cash_balance"],
             total_money_in=result["total_money_in"],
             holdings=[
-                Holding(symbol=holding["symbol"], shares=holding["shares"], total_cost=holding["total_cost"])
+                Holding(
+                    symbol=holding["symbol"],
+                    shares=holding["shares"],
+                    stock_type=StockType(holding["stock_type"]),
+                    total_cost=holding["total_cost"],
+                )
                 for holding in result["holdings"]
             ],
             created_at=result["created_at"],
@@ -31,7 +36,7 @@ def get(self, user_id: int) -> Portfolio:
 
     def update(self, portfolio: Portfolio) -> None:
         portfolio.updated_at = datetime.now(timezone.utc)
-        self.collection.replace_one({"user_id": portfolio.user_id}, asdict(portfolio), upsert=True)
+        self.collection.replace_one({"user_id": portfolio.user_id}, portfolio.as_dict(), upsert=True)
 
     def __del__(self):
         self.client.close()

src/adapters/stock.py

@@ -1,10 +1,9 @@
 from typing import List
-
 from pymongo import MongoClient
 from pymongo.database import Database
-
 from .base import AbstractStockRepository
-from domain.stock import CreateStock, Stock, ActionType
+from domain.stock import CreateStock, Stock
+from domain.enum import ActionType, StockType
 
 
 class StockRepository(AbstractStockRepository):
@@ -14,17 +13,7 @@ def __init__(self, mongo_client: MongoClient, database_name: str = "stock_db"):
         self.collection = self.db["stocks"]
 
     def create(self, stock: CreateStock) -> str:
-        stock_dict = {
-            "user_id": stock.user_id,
-            "symbol": stock.symbol,
-            "price": stock.price,
-            "quantity": stock.quantity,
-            "action_type": stock.action_type.value,
-            "created_at": stock.created_at,
-            "updated_at": stock.created_at,
-        }
-
-        result = self.collection.insert_one(stock_dict)
+        result = self.collection.insert_one(stock.as_dict())
         return str(result.inserted_id)
 
     def list(self, user_id: int) -> List[Stock]:
@@ -38,6 +27,7 @@ def list(self, user_id: int) -> List[Stock]:
                 price=doc["price"],
                 quantity=doc["quantity"],
                 action_type=ActionType(doc["action_type"]),
+                stock_type=StockType(doc["stock_type"]),
                 created_at=doc["created_at"],
                 updated_at=doc["updated_at"],
             )

src/domain/enum.py

@@ -0,0 +1,25 @@
+from enum import Enum
+
+
+class ActionType(Enum):
+    BUY = "BUY"
+    SELL = "SELL"
+    TRANSFER = "TRANSFER"
+
+
+ACTION_MAP = {
+    1: ActionType.BUY,
+    2: ActionType.SELL,
+    3: ActionType.TRANSFER,
+}
+
+
+class StockType(Enum):
+    STOCKS = "STOCKS"
+    ETF = "ETF"
+
+
+STOCK_MAP = {
+    1: StockType.STOCKS,
+    2: StockType.ETF,
+}

src/domain/portfolio.py

@@ -1,14 +1,44 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, asdict
 from datetime import datetime
-from typing import List
+from typing import List, TypedDict
+from .enum import StockType
+from utils.utils import custom_dict_factory
+
+
+class HoldingDict(TypedDict):
+    symbol: str
+    shares: int
+    stock_type: str
+    total_cost: float
+
+
+class PortfolioDict(TypedDict):
+    user_id: int
+    cash_balance: float
+    total_money_in: float
+    holdings: List[HoldingDict]
+    created_at: datetime
+    updated_at: datetime
 
 
 @dataclass
 class Holding:
     symbol: str
     shares: int
+    stock_type: StockType
     total_cost: float
 
+    def __post_init__(self):
+        if not self.symbol:
+            raise ValueError("symbol cannot be empty")
+        if self.shares < 0:
+            raise ValueError("shares cannot be negative")
+        if self.total_cost < 0:
+            raise ValueError("total_cost cannot be negative")
+
+    def as_dict(self) -> HoldingDict:
+        return asdict(self, dict_factory=custom_dict_factory)
+
 
 @dataclass
 class Portfolio:
@@ -18,3 +48,12 @@ class Portfolio:
     holdings: List[Holding]
     created_at: datetime
     updated_at: datetime
+
+    def __post_init__(self):
+        if self.cash_balance < 0:
+            raise ValueError("cash_balance cannot be negative")
+        if self.total_money_in < 0:
+            raise ValueError("total_money_in cannot be negative")
+
+    def as_dict(self) -> PortfolioDict:
+        return asdict(self, dict_factory=custom_dict_factory)

src/domain/stock.py

@@ -1,15 +1,30 @@
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
+from typing import TypedDict
 from datetime import datetime
-from enum import Enum
+from utils.utils import custom_dict_factory
+from .enum import ActionType, StockType
 
 
-class ActionType(Enum):
-    BUY = "BUY"
-    SELL = "SELL"
-    TRANSFER = "TRANSFER"
+class CreateStockDict(TypedDict):
+    user_id: int
+    symbol: str
+    price: float
+    quantity: int
+    action_type: str
+    stock_type: str
+    created_at: datetime
 
 
-ACTION_MAP = {1: ActionType.BUY, 2: ActionType.SELL, 3: ActionType.TRANSFER}
+class StockDict(TypedDict):
+    id: str
+    user_id: int
+    symbol: str
+    price: float
+    quantity: int
+    action_type: str
+    stock_type: str
+    created_at: datetime
+    updated_at: datetime
 
 
 @dataclass
@@ -19,8 +34,22 @@ class CreateStock:
     price: float
     quantity: int
     action_type: ActionType
+    stock_type: StockType
     created_at: datetime
 
+    def __post_init__(self):
+        if not self.user_id or self.user_id <= 0:
+            raise ValueError("user_id must be non-empty and greater than 0")
+        if not self.symbol or self.symbol.strip() == "":
+            raise ValueError("symbol must be a non-empty string")
+        if self.price <= 0:
+            raise ValueError("price must be greater than 0")
+        if self.quantity <= 0:
+            raise ValueError("quantity must be greater than 0")
+
+    def as_dict(self) -> CreateStockDict:
+        return asdict(self, dict_factory=custom_dict_factory)
+
 
 @dataclass
 class Stock:
@@ -30,5 +59,21 @@ class Stock:
     price: float
     quantity: int
     action_type: ActionType
+    stock_type: StockType
     created_at: datetime
     updated_at: datetime
+
+    def __post_init__(self):
+        if not self.id:
+            raise ValueError("id cannot be empty")
+        if self.user_id <= 0:
+            raise ValueError("user_id must be positive")
+        if not self.symbol:
+            raise ValueError("symbol cannot be empty")
+        if self.price < 0:
+            raise ValueError("price cannot be negative")
+        if self.quantity <= 0:
+            raise ValueError("quantity must be positive")
+
+    def as_dict(self) -> StockDict:
+        return asdict(self, dict_factory=custom_dict_factory)

src/handler/stock.py

@@ -1,12 +1,11 @@
 import logging
 from datetime import datetime, timezone
-
 import grpc
 import proto.stock_pb2 as stock_pb2
 import proto.stock_pb2_grpc as stock_pb2_grpc
-
 from usecase.base import AbstractStockUsecase
-from domain.stock import CreateStock, ActionType, ACTION_MAP
+from domain.stock import CreateStock
+from domain.enum import ActionType, ACTION_MAP, StockType, STOCK_MAP
 
 
 class StockService(stock_pb2_grpc.StockService):
@@ -15,13 +14,13 @@ def __init__(self, stock_usecase: AbstractStockUsecase):
 
     def Create(self, request, context):
         try:
-            self._validate_create_request(request)
             stock = CreateStock(
                 user_id=request.user_id,
                 symbol=request.symbol,
                 price=request.price,
                 quantity=request.quantity,
                 action_type=self._map_action_type(request.action),
+                stock_type=self._map_stock_type(request.stock_type),
                 created_at=datetime.now(timezone.utc),
             )
 
@@ -64,15 +63,10 @@ def _map_action_type(self, action: int) -> ActionType:
             raise ValueError(f"Invalid action type: {action}. Must be 1 (BUY), 2 (SELL), or 3 (TRANSFER).")
         return ACTION_MAP[action]
 
-    def _validate_create_request(self, request):
-        if not request.user_id or request.user_id <= 0:
-            raise ValueError("user_id must be non-empty and greater than 0")
-        if not request.symbol or request.symbol.strip() == "":
-            raise ValueError("symbol must be a non-empty string")
-        if request.price <= 0:
-            raise ValueError("price must be greater than 0")
-        if request.quantity <= 0:
-            raise ValueError("quantity must be greater than 0")
+    def _map_stock_type(self, stock_type: int) -> StockType:
+        if stock_type not in STOCK_MAP:
+            raise ValueError(f"Invalid stock type: {stock_type}. Must be 1 (STOCKS), 2 (ETF).")
+        return STOCK_MAP[stock_type]
 
     def _convert_to_proto_stock_list(self, stock_list):
         return [
@@ -83,6 +77,7 @@ def _convert_to_proto_stock_list(self, stock_list):
                 price=stock.price,
                 quantity=stock.quantity,
                 action=stock.action_type.value,
+                stock_type=stock.stock_type.value,
                 created_at=stock.created_at,
                 updated_at=stock.updated_at,
             )

src/proto/stock.proto

@@ -13,15 +13,24 @@ message Action {
   }
 }
 
+message StockType {
+  enum Type {
+    UNSPECIFIED = 0;
+    STOCKS = 1;
+    ETF = 2;
+  }
+}
+
 message Stock {
   string id = 1 [json_name = "id"];
   int32 user_id = 2 [json_name = "user_id"];
   string symbol = 3 [json_name = "symbol"];
   double price = 4 [json_name = "price"];
   int32 quantity = 5 [json_name = "quantity"];
   string action = 6 [json_name = "action"];
-  google.protobuf.Timestamp created_at = 7 [json_name = "created_at"];
-  google.protobuf.Timestamp updated_at = 8 [json_name = "updated_at"];
+  string stock_type = 7 [json_name = "stock_type"];
+  google.protobuf.Timestamp created_at = 8 [json_name = "created_at"];
+  google.protobuf.Timestamp updated_at = 9 [json_name = "updated_at"];
 }
 
 message CreateReq {
@@ -30,8 +39,9 @@ message CreateReq {
   double price = 3 [json_name = "price"];
   int32 quantity = 4 [json_name = "quantity"];
   Action.Type action = 5 [json_name = "action"]; // add validation rules
-  google.protobuf.Timestamp created_at = 6 [json_name = "created_at"];
-  google.protobuf.Timestamp updated_at = 7 [json_name = "updated_at"];
+  StockType.Type stock_type = 6 [json_name = "stock_type"]; // add validation rules
+  google.protobuf.Timestamp created_at = 7 [json_name = "created_at"];
+  google.protobuf.Timestamp updated_at = 8 [json_name = "updated_at"];
 }
 
 message CreateResp {