From ed70736fb489673909d226b5d81ce6d0147618ba Mon Sep 17 00:00:00 2001
From: Kun Jinkao <45487685+Snoopy1866@users.noreply.github.com>
Date: Fri, 13 Sep 2024 16:56:55 +0800
Subject: [PATCH] Update two_proportion.py

---
 src/pystatpower/procedures/two_proportion.py | 35 +++++++++++---------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/src/pystatpower/procedures/two_proportion.py b/src/pystatpower/procedures/two_proportion.py
index b5c11a1..95aacde 100644
--- a/src/pystatpower/procedures/two_proportion.py
+++ b/src/pystatpower/procedures/two_proportion.py
@@ -3,6 +3,16 @@
 from enum import Enum
 
 
+class SolvableParameter(Enum):
+    """求解目标"""
+
+    N = 1
+    ALPHA = 2
+    POWER = 3
+    TREATMENT_PROPORTION = 4
+    REFERENCE_PROPORTION = 5
+
+
 class Alternative(Enum):
     """假设检验的备择假设类型"""
 
@@ -32,25 +42,20 @@ class GroupAllocation(Enum):
 
 
 class TwoProportionDesigner:
-    def __init__(self, solve_for: str):
-        if not isinstance(solve_for, str):
-            raise TypeError("solve_for must be a string")
-
-        self._solve_for = solve_for.lower()
-        if self._solve_for not in ["n", "alpha", "power", "treatment_proportion", "reference_proportion"]:
-            raise ValueError(
-                "solve_for must be either 'n', 'alpha', 'power', 'treatment_proportion', or 'reference_proportion'"
-            )
-        match self._solve_for:
-            case "n":
+    def __init__(self, solve_for: SolvableParameter):
+        if not isinstance(solve_for, SolvableParameter):
+            raise TypeError("solve_for must be an instance of SolvableParameter")
+
+        match solve_for:
+            case SolvableParameter.N:
                 return TwoProportionSolveForNDesigner()
-            case "alpha":
+            case SolvableParameter.ALPHA:
                 return TwoProportionSolveForAlphaDesigner()
-            case "power":
+            case SolvableParameter.POWER:
                 return TwoProportionSolveForPowerDesigner()
-            case "treatment_proportion":
+            case SolvableParameter.TREATMENT_PROPORTION:
                 return TwoProportionSolveForTreatmentProportionDesigner()
-            case "reference_proportion":
+            case SolvableParameter.REFERENCE_PROPORTION:
                 return TwoProportionSolveForReferenceProportionDesigner()