fix: PR comments

limebit · Jan 21, 2025 · 2a405db · 2a405db
1 parent 9ef8d28
commit 2a405db
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/medmodels/treatment_effect/matching/neighbors.py b/medmodels/treatment_effect/matching/neighbors.py
@@ -55,10 +55,10 @@ def match_controls(
             control_set (Set[NodeIndex]): Set of control subjects.
             treated_set (Set[NodeIndex]): Set of treated subjects.
             patients_group (Group): Group of patients in the MedRecord.
-            essential_covariates (Optional[List[MedRecordAttribute]], optional):
+            essential_covariates (Optional[Sequence[MedRecordAttribute]], optional):
                 Covariates that are essential for matching. Defaults to None, meaning
                 all the attributes of the patients are used.
-            one_hot_covariates (Optional[List[MedRecordAttribute]], optional):
+            one_hot_covariates (Optional[Sequence[MedRecordAttribute]], optional):
                 Covariates that are one-hot encoded for matching. Defaults to None,
                 meaning all the categorical attributes of the patients are used.
 

diff --git a/medmodels/treatment_effect/matching/propensity.py b/medmodels/treatment_effect/matching/propensity.py
@@ -75,10 +75,10 @@ def match_controls(
             control_set (Set[NodeIndex]): Set of control subjects.
             treated_set (Set[NodeIndex]): Set of treated subjects.
             patients_group (Group): Group of patients in MedRecord.
-            essential_covariates (Optional[List[MedRecordAttribute]], optional):
+            essential_covariates (Optional[Sequence[MedRecordAttribute]], optional):
                 Covariates that are essential for matching. Defaults to None, meaning
                 all the attributes of the patients are used.
-            one_hot_covariates (Optional[List[MedRecordAttribute]], optional):
+            one_hot_covariates (Optional[Sequence[MedRecordAttribute]], optional):
                 Covariates that are one-hot encoded for matching. Defaults to None,
                 meaning all the categorical attributes of the patients are used.
 

diff --git a/medmodels/treatment_effect/treatment_effect.py b/medmodels/treatment_effect/treatment_effect.py
@@ -68,8 +68,8 @@ class TreatmentEffect:
     _filter_controls_query: Optional[NodeQuery]
 
     _matching_method: Optional[MatchingMethod]
-    _matching_essential_covariates: Optional[MedRecordAttributeInputList]
-    _matching_one_hot_covariates: Optional[MedRecordAttributeInputList]
+    _matching_essential_covariates: MedRecordAttributeInputList
+    _matching_one_hot_covariates: MedRecordAttributeInputList
     _matching_model: Model
     _matching_number_of_neighbors: int
     _matching_hyperparameters: Optional[Dict[str, Any]]