Fix pydantic issue and update variable names to more meaningful names…

… in models.py
rgaveiga · Sep 27, 2024 · ec03b85 · ec03b85
1 parent 3e6e201
commit ec03b85
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 26 deletions.
diff --git a/optionlab/__init__.py b/optionlab/__init__.py
@@ -1,7 +1,7 @@
 import typing
 
 
-VERSION = "1.3.0"
+VERSION = "1.3.1"
 
 
 if typing.TYPE_CHECKING:
@@ -11,16 +11,16 @@
         Inputs,
         OptionType,
         OptionInfo,
-        OptionStrategy,
+        Option,
         Outputs,
         ClosedPosition,
         ProbabilityOfProfitArrayInputs,
         ProbabilityOfProfitInputs,
         BlackScholesInfo,
         Distribution,
-        Strategy,
+        StrategyLeg,
         StrategyType,
-        StockStrategy,
+        Stock,
         Country,
         Action,
     )
@@ -53,16 +53,16 @@
     "Inputs",
     "OptionType",
     "OptionInfo",
-    "OptionStrategy",
+    "Option",
     "Outputs",
     "ClosedPosition",
     "ProbabilityOfProfitArrayInputs",
     "ProbabilityOfProfitInputs",
     "BlackScholesInfo",
     "Distribution",
-    "Strategy",
+    "StrategyLeg",
     "StrategyType",
-    "StockStrategy",
+    "Stock",
     "Country",
     "Action",
     # engine
@@ -96,15 +96,15 @@
     "Outputs": (__package__, ".models"),
     "OptionType": (__package__, ".models"),
     "OptionInfo": (__package__, ".models"),
-    "OptionStrategy": (__package__, ".models"),
+    "Option": (__package__, ".models"),
     "ClosedPosition": (__package__, ".models"),
     "ProbabilityOfProfitArrayInputs": (__package__, ".models"),
     "ProbabilityOfProfitInputs": (__package__, ".models"),
     "BlackScholesInfo": (__package__, ".models"),
     "Distribution": (__package__, ".models"),
-    "Strategy": (__package__, ".models"),
+    "StrategyLeg": (__package__, ".models"),
     "StrategyType": (__package__, ".models"),
-    "StockStrategy": (__package__, ".models"),
+    "Stock": (__package__, ".models"),
     "Country": (__package__, ".models"),
     "Action": (__package__, ".models"),
     # engine

diff --git a/optionlab/engine.py b/optionlab/engine.py
@@ -12,8 +12,8 @@
 from optionlab.models import (
     Inputs,
     Action,
-    OptionStrategy,
-    StockStrategy,
+    Option,
+    Stock,
     ClosedPosition,
     Outputs,
     ProbabilityOfProfitInputs,
@@ -75,7 +75,7 @@ def _init_inputs(inputs: Inputs) -> EngineData:
     for i, strategy in enumerate(inputs.strategy):
         data.type.append(strategy.type)
 
-        if isinstance(strategy, OptionStrategy):
+        if isinstance(strategy, Option):
             data.strike.append(strategy.strike)
             data.premium.append(strategy.premium)
             data.n.append(strategy.n)
@@ -110,7 +110,7 @@ def _init_inputs(inputs: Inputs) -> EngineData:
             else:
                 raise ValueError("Expiration must be a date, an int or None.")
 
-        elif isinstance(strategy, StockStrategy):
+        elif isinstance(strategy, Stock):
             data.n.append(strategy.n)
             data.action.append(strategy.action)
             data._previous_position.append(strategy.prev_pos or 0.0)

diff --git a/optionlab/models.py b/optionlab/models.py
@@ -27,12 +27,13 @@
 ]
 
 
-class BaseStrategy(BaseModel):
+class BaseLeg(BaseModel):
+    n: int = Field(gt=0)
     action: Action
     prev_pos: float | None = None
 
 
-class StockStrategy(BaseStrategy):
+class Stock(BaseLeg):
     """
     "type" : string
         It must be 'stock'. It is mandatory.
@@ -48,15 +49,12 @@ class StockStrategy(BaseStrategy):
         negative, it means that the position is closed and the
         difference between this price and the current price is
         considered in the payoff calculation.
-
     """
 
     type: Literal["stock"] = "stock"
-    n: int = Field(gt=0)
-    premium: float | None = None
 
 
-class OptionStrategy(BaseStrategy):
+class Option(BaseLeg):
     """
     "type" : string
         Either 'call' or 'put'. It is mandatory.
@@ -83,7 +81,6 @@ class OptionStrategy(BaseStrategy):
     type: OptionType
     strike: float = Field(gt=0)
     premium: float = Field(gt=0)
-    n: int = Field(gt=0)
     expiration: dt.date | int | None = None
 
     @field_validator("expiration")
@@ -107,7 +104,7 @@ class ClosedPosition(BaseModel):
     prev_pos: float
 
 
-Strategy = StockStrategy | OptionStrategy | ClosedPosition
+StrategyLeg = Stock | Option | ClosedPosition
 
 
 class ProbabilityOfProfitInputs(BaseModel):
@@ -203,7 +200,7 @@ class Inputs(BaseModel):
     interest_rate: float = Field(gt=0, le=0.2)
     min_stock: float
     max_stock: float
-    strategy: list[Strategy] = Field(..., min_length=1, discriminator="type")
+    strategy: list[StrategyLeg] = Field(..., min_length=1)
     dividend_yield: float = 0.0
     profit_target: float | None = None
     loss_limit: float | None = None
@@ -221,7 +218,7 @@ class Inputs(BaseModel):
 
     @field_validator("strategy")
     @classmethod
-    def validate_strategy(cls, v: list[Strategy]) -> list[Strategy]:
+    def validate_strategy(cls, v: list[StrategyLeg]) -> list[StrategyLeg]:
         types = [strategy.type for strategy in v]
         if types.count("closed") > 1:
             raise ValueError("Only one position of type 'closed' is allowed!")
@@ -232,8 +229,7 @@ def validate_dates(self) -> "Inputs":
         expiration_dates = [
             strategy.expiration
             for strategy in self.strategy
-            if isinstance(strategy, OptionStrategy)
-            and isinstance(strategy.expiration, dt.date)
+            if isinstance(strategy, Option) and isinstance(strategy.expiration, dt.date)
         ]
         if self.start_date and self.target_date:
             if any(