Scenarios keep method

edsl/scenarios/Scenario.py

@@ -216,6 +216,23 @@ def drop(self, list_of_keys: List[str]) -> "Scenario":
                 new_scenario[key] = self[key]
         return new_scenario
 
+    def keep(self, list_of_keys: List[str]) -> "Scenario":
+        """Keep a subset of keys from a scenario.
+
+        :param list_of_keys: The keys to keep.
+
+        Example:
+
+        >>> s = Scenario({"food": "wood chips", "drink": "water"})
+        >>> s.keep(["food"])
+        Scenario({'food': 'wood chips'})
+        """
+        new_scenario = Scenario()
+        for key in self.keys():
+            if key in list_of_keys:
+                new_scenario[key] = self[key]
+        return new_scenario
+
     @classmethod
     def from_url(cls, url: str, field_name: Optional[str] = "text") -> "Scenario":
         """Creates a scenario from a URL.

edsl/scenarios/ScenarioList.py

@@ -538,6 +538,17 @@ def drop(self, *fields) -> ScenarioList:
         """
         return ScenarioList([scenario.drop(fields) for scenario in self.data])
 
+    def keep(self, *fields) -> ScenarioList:
+        """Keep only the specified fields in the scenarios.
+
+        Example:
+
+        >>> s = ScenarioList([Scenario({'a': 1, 'b': 1}), Scenario({'a': 1, 'b': 2})])
+        >>> s.keep('a')
+        ScenarioList([Scenario({'a': 1}), Scenario({'a': 1})])
+        """
+        return ScenarioList([scenario.keep(fields) for scenario in self.data])
+
     @classmethod
     def from_list(
         cls, name: str, values: list, func: Optional[Callable] = None