edsl 0.1.33.dev1__py3-none-any.whl → 0.1.33.dev2__py3-none-any.whl

edsl/scenarios/ScenarioList.py

@@ -6,6 +6,11 @@ import csv
 import random
 from collections import UserList, Counter
 from collections.abc import Iterable
+import urllib.parse
+import urllib.request
+from io import StringIO
+from collections import defaultdict
+import inspect
 
 from simpleeval import EvalWithCompoundTypes
 
@@ -15,6 +20,9 @@ from edsl.scenarios.Scenario import Scenario
 from edsl.scenarios.ScenarioListPdfMixin import ScenarioListPdfMixin
 from edsl.scenarios.ScenarioListExportMixin import ScenarioListExportMixin
 
+from edsl.conjure.naming_utilities import sanitize_string
+from edsl.utilities.utilities import is_valid_variable_name
+
 
 class ScenarioListMixin(ScenarioListPdfMixin, ScenarioListExportMixin):
     pass
@@ -23,12 +31,171 @@ class ScenarioListMixin(ScenarioListPdfMixin, ScenarioListExportMixin):
 class ScenarioList(Base, UserList, ScenarioListMixin):
     """Class for creating a list of scenarios to be used in a survey."""
 
-    def __init__(self, data: Optional[list] = None):
+    def __init__(self, data: Optional[list] = None, codebook: Optional[dict] = None):
         """Initialize the ScenarioList class."""
         if data is not None:
             super().__init__(data)
         else:
             super().__init__([])
+        self.codebook = codebook or {}
+
+    def give_valid_names(self) -> ScenarioList:
+        """Give valid names to the scenario keys.
+
+        >>> s = ScenarioList([Scenario({'a': 1, 'b': 2}), Scenario({'a': 1, 'b': 1})])
+        >>> s.give_valid_names()
+        ScenarioList([Scenario({'a': 1, 'b': 2}), Scenario({'a': 1, 'b': 1})])
+        >>> s = ScenarioList([Scenario({'are you there John?': 1, 'b': 2}), Scenario({'a': 1, 'b': 1})])
+        >>> s.give_valid_names()
+        ScenarioList([Scenario({'john': 1, 'b': 2}), Scenario({'a': 1, 'b': 1})])
+        """
+        codebook = {}
+        new_scenaerios = []
+        for scenario in self:
+            new_scenario = {}
+            for key in scenario:
+                if not is_valid_variable_name(key):
+                    if key in codebook:
+                        new_key = codebook[key]
+                    else:
+                        new_key = sanitize_string(key)
+                        if not is_valid_variable_name(new_key):
+                            new_key = f"var_{len(codebook)}"
+                        codebook[key] = new_key
+                    new_scenario[new_key] = scenario[key]
+                else:
+                    new_scenario[key] = scenario[key]
+            new_scenaerios.append(Scenario(new_scenario))
+        return ScenarioList(new_scenaerios, codebook)
+
+    def unpivot(self, id_vars=None, value_vars=None):
+        """
+        Unpivot the ScenarioList, allowing for id variables to be specified.
+
+        Parameters:
+        id_vars (list): Fields to use as identifier variables (kept in each entry)
+        value_vars (list): Fields to unpivot. If None, all fields not in id_vars will be used.
+
+        Example:
+        >>> s = ScenarioList([
+        ...     Scenario({'id': 1, 'year': 2020, 'a': 10, 'b': 20}),
+        ...     Scenario({'id': 2, 'year': 2021, 'a': 15, 'b': 25})
+        ... ])
+        >>> s.unpivot(id_vars=['id', 'year'], value_vars=['a', 'b'])
+        ScenarioList([Scenario({'id': 1, 'year': 2020, 'variable': 'a', 'value': 10}), Scenario({'id': 1, 'year': 2020, 'variable': 'b', 'value': 20}), Scenario({'id': 2, 'year': 2021, 'variable': 'a', 'value': 15}), Scenario({'id': 2, 'year': 2021, 'variable': 'b', 'value': 25})])
+        """
+        if id_vars is None:
+            id_vars = []
+        if value_vars is None:
+            value_vars = [field for field in self[0].keys() if field not in id_vars]
+
+        new_scenarios = []
+        for scenario in self:
+            for var in value_vars:
+                new_scenario = {id_var: scenario[id_var] for id_var in id_vars}
+                new_scenario["variable"] = var
+                new_scenario["value"] = scenario[var]
+                new_scenarios.append(Scenario(new_scenario))
+
+        return ScenarioList(new_scenarios)
+
+    def pivot(self, id_vars, var_name="variable", value_name="value"):
+        """
+        Pivot the ScenarioList from long to wide format.
+
+        Parameters:
+        id_vars (list): Fields to use as identifier variables
+        var_name (str): Name of the variable column (default: 'variable')
+        value_name (str): Name of the value column (default: 'value')
+
+        Example:
+        >>> s = ScenarioList([
+        ...     Scenario({'id': 1, 'year': 2020, 'variable': 'a', 'value': 10}),
+        ...     Scenario({'id': 1, 'year': 2020, 'variable': 'b', 'value': 20}),
+        ...     Scenario({'id': 2, 'year': 2021, 'variable': 'a', 'value': 15}),
+        ...     Scenario({'id': 2, 'year': 2021, 'variable': 'b', 'value': 25})
+        ... ])
+        >>> s.pivot(id_vars=['id', 'year'])
+        ScenarioList([Scenario({'id': 1, 'year': 2020, 'a': 10, 'b': 20}), Scenario({'id': 2, 'year': 2021, 'a': 15, 'b': 25})])
+        """
+        pivoted_dict = {}
+
+        for scenario in self:
+            # Create a tuple of id values to use as a key
+            id_key = tuple(scenario[id_var] for id_var in id_vars)
+
+            # If this combination of id values hasn't been seen before, initialize it
+            if id_key not in pivoted_dict:
+                pivoted_dict[id_key] = {id_var: scenario[id_var] for id_var in id_vars}
+
+            # Add the variable-value pair to the dict
+            variable = scenario[var_name]
+            value = scenario[value_name]
+            pivoted_dict[id_key][variable] = value
+
+        # Convert the dict of dicts to a list of Scenarios
+        pivoted_scenarios = [
+            Scenario(dict(zip(id_vars, id_key), **values))
+            for id_key, values in pivoted_dict.items()
+        ]
+
+        return ScenarioList(pivoted_scenarios)
+
+    def group_by(self, id_vars, variables, func):
+        """
+        Group the ScenarioList by id_vars and apply a function to the specified variables.
+
+        Parameters:
+        id_vars (list): Fields to use as identifier variables for grouping
+        variables (list): Fields to pass to the aggregation function
+        func (callable): Function to apply to the grouped variables.
+                        Should accept lists of values for each variable.
+
+        Returns:
+        ScenarioList: A new ScenarioList with the grouped and aggregated results
+
+        Example:
+        >>> def avg_sum(a, b):
+        ...     return {'avg_a': sum(a) / len(a), 'sum_b': sum(b)}
+        >>> s = ScenarioList([
+        ...     Scenario({'group': 'A', 'year': 2020, 'a': 10, 'b': 20}),
+        ...     Scenario({'group': 'A', 'year': 2021, 'a': 15, 'b': 25}),
+        ...     Scenario({'group': 'B', 'year': 2020, 'a': 12, 'b': 22}),
+        ...     Scenario({'group': 'B', 'year': 2021, 'a': 17, 'b': 27})
+        ... ])
+        >>> s.group_by(id_vars=['group'], variables=['a', 'b'], func=avg_sum)
+        ScenarioList([Scenario({'group': 'A', 'avg_a': 12.5, 'sum_b': 45}), Scenario({'group': 'B', 'avg_a': 14.5, 'sum_b': 49})])
+        """
+        # Check if the function is compatible with the specified variables
+        func_params = inspect.signature(func).parameters
+        if len(func_params) != len(variables):
+            raise ValueError(
+                f"Function {func.__name__} expects {len(func_params)} arguments, but {len(variables)} variables were provided"
+            )
+
+        # Group the scenarios
+        grouped = defaultdict(lambda: defaultdict(list))
+        for scenario in self:
+            key = tuple(scenario[id_var] for id_var in id_vars)
+            for var in variables:
+                grouped[key][var].append(scenario[var])
+
+        # Apply the function to each group
+        result = []
+        for key, group in grouped.items():
+            try:
+                aggregated = func(*[group[var] for var in variables])
+            except Exception as e:
+                raise ValueError(f"Error applying function to group {key}: {str(e)}")
+
+            if not isinstance(aggregated, dict):
+                raise ValueError(f"Function {func.__name__} must return a dictionary")
+
+            new_scenario = dict(zip(id_vars, key))
+            new_scenario.update(aggregated)
+            result.append(Scenario(new_scenario))
+
+        return ScenarioList(result)
 
     @property
     def parameters(self) -> set:
@@ -153,6 +320,71 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
                 new_scenarios.append(new_scenario)
         return ScenarioList(new_scenarios)
 
+    def concatenate(self, fields: List[str], separator: str = ";") -> "ScenarioList":
+        """Concatenate specified fields into a single field.
+
+        Args:
+            fields (List[str]): List of field names to concatenate.
+            separator (str, optional): Separator to use between field values. Defaults to ";".
+
+        Returns:
+            ScenarioList: A new ScenarioList with concatenated fields.
+
+        Example:
+            >>> s = ScenarioList([Scenario({'a': 1, 'b': 2, 'c': 3}), Scenario({'a': 4, 'b': 5, 'c': 6})])
+            >>> s.concatenate(['a', 'b', 'c'])
+            ScenarioList([Scenario({'concat_a_b_c': '1;2;3'}), Scenario({'concat_a_b_c': '4;5;6'})])
+        """
+        new_scenarios = []
+        for scenario in self:
+            new_scenario = scenario.copy()
+            concat_values = []
+            for field in fields:
+                if field in new_scenario:
+                    concat_values.append(str(new_scenario[field]))
+                    del new_scenario[field]
+
+            new_field_name = f"concat_{'_'.join(fields)}"
+            new_scenario[new_field_name] = separator.join(concat_values)
+            new_scenarios.append(new_scenario)
+
+        return ScenarioList(new_scenarios)
+
+    def unpack_dict(
+        self, field: str, prefix: Optional[str] = None, drop_field: bool = False
+    ) -> ScenarioList:
+        """Unpack a dictionary field into separate fields.
+
+        Example:
+
+        >>> s = ScenarioList([Scenario({'a': 1, 'b': {'c': 2, 'd': 3}})])
+        >>> s.unpack_dict('b')
+        ScenarioList([Scenario({'a': 1, 'b': {'c': 2, 'd': 3}, 'c': 2, 'd': 3})])
+        """
+        new_scenarios = []
+        for scenario in self:
+            new_scenario = scenario.copy()
+            for key, value in scenario[field].items():
+                if prefix:
+                    new_scenario[prefix + key] = value
+                else:
+                    new_scenario[key] = value
+            if drop_field:
+                new_scenario.pop(field)
+            new_scenarios.append(new_scenario)
+        return ScenarioList(new_scenarios)
+
+    def transform(
+        self, field: str, func: Callable, new_name: Optional[str] = None
+    ) -> ScenarioList:
+        """Transform a field using a function."""
+        new_scenarios = []
+        for scenario in self:
+            new_scenario = scenario.copy()
+            new_scenario[new_name or field] = func(scenario[field])
+            new_scenarios.append(new_scenario)
+        return ScenarioList(new_scenarios)
+
     def mutate(
         self, new_var_string: str, functions_dict: Optional[dict[str, Callable]] = None
     ) -> ScenarioList:
@@ -310,6 +542,19 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
         data = [{key: [scenario[key] for scenario in self.data]} for key in keys]
         return Dataset(data)
 
+    def split(
+        self, field: str, split_on: str, index: int, new_name: Optional[str] = None
+    ) -> ScenarioList:
+        """Split a scenario fiel in multiple fields."""
+        if new_name is None:
+            new_name = field + "_split_" + str(index)
+        new_scenarios = []
+        for scenario in self:
+            new_scenario = scenario.copy()
+            new_scenario[new_name] = scenario[field].split(split_on)[index]
+            new_scenarios.append(new_scenario)
+        return ScenarioList(new_scenarios)
+
     def add_list(self, name, values) -> ScenarioList:
         """Add a list of values to a ScenarioList.
 
@@ -392,37 +637,6 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
 
         return ScenarioList([Scenario(entry) for entry in processed_lines])
 
-    @classmethod
-    def from_docx(cls, docx_file_path: str):
-        from docx import Document
-
-        doc = Document(docx_file_path)
-        lines = []
-
-        # Extract text from paragraphs, treating each paragraph as a line
-        for para in doc.paragraphs:
-            lines.extend(para.text.splitlines())
-
-        processed_lines = []
-        non_blank_lines = [
-            (i, line.strip()) for i, line in enumerate(lines) if line.strip()
-        ]
-
-        for index, (line_no, text) in enumerate(non_blank_lines):
-            entry = {
-                "line_no": line_no + 1,  # Using 1-based index for line numbers
-                "text": text,
-                "line_before": non_blank_lines[index - 1][1] if index > 0 else None,
-                "line_after": (
-                    non_blank_lines[index + 1][1]
-                    if index < len(non_blank_lines) - 1
-                    else None
-                ),
-            }
-            processed_lines.append(entry)
-
-        return ScenarioList([Scenario(entry) for entry in processed_lines])
-
     @classmethod
     def from_google_doc(cls, url: str) -> ScenarioList:
         """Create a ScenarioList from a Google Doc.
@@ -473,6 +687,62 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
         """
         return cls([Scenario(row) for row in df.to_dict(orient="records")])
 
+    @classmethod
+    def from_wikipedia(cls, url: str, table_index: int = 0):
+        """
+        Extracts a table from a Wikipedia page.
+
+        Parameters:
+            url (str): The URL of the Wikipedia page.
+            table_index (int): The index of the table to extract (default is 0).
+
+        Returns:
+            pd.DataFrame: A DataFrame containing the extracted table.
+        # # Example usage
+        # url = "https://en.wikipedia.org/wiki/List_of_countries_by_GDP_(nominal)"
+        # df = from_wikipedia(url, 0)
+
+        # if not df.empty:
+        #     print(df.head())
+        # else:
+        #     print("Failed to extract table.")
+
+
+        """
+        import pandas as pd
+        import requests
+        from requests.exceptions import RequestException
+
+        try:
+            # Check if the URL is reachable
+            response = requests.get(url)
+            response.raise_for_status()  # Raises HTTPError for bad responses
+
+            # Extract tables from the Wikipedia page
+            tables = pd.read_html(url)
+
+            # Ensure the requested table index is within the range of available tables
+            if table_index >= len(tables) or table_index < 0:
+                raise IndexError(
+                    f"Table index {table_index} is out of range. This page has {len(tables)} table(s)."
+                )
+
+            # Return the requested table as a DataFrame
+            # return tables[table_index]
+            return cls.from_pandas(tables[table_index])
+
+        except RequestException as e:
+            print(f"Error fetching the URL: {e}")
+        except ValueError as e:
+            print(f"Error parsing tables: {e}")
+        except IndexError as e:
+            print(e)
+        except Exception as e:
+            print(f"An unexpected error occurred: {e}")
+
+        # Return an empty DataFrame in case of an error
+        # return cls.from_pandas(pd.DataFrame())
+
     def to_key_value(self, field: str, value=None) -> Union[dict, set]:
         """Return the set of values in the field.
 
@@ -594,8 +864,15 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
         return cls.from_excel(temp_filename, sheet_name=sheet_name)
 
     @classmethod
-    def from_csv(cls, filename: str) -> ScenarioList:
-        """Create a ScenarioList from a CSV file.
+    def from_csv(cls, source: Union[str, urllib.parse.ParseResult]) -> ScenarioList:
+        """Create a ScenarioList from a CSV file or URL.
+
+        Args:
+            source: A string representing either a local file path or a URL to a CSV file,
+                    or a urllib.parse.ParseResult object for a URL.
+
+        Returns:
+            ScenarioList: A ScenarioList object containing the data from the CSV.
 
         Example:
 
@@ -611,15 +888,37 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
         'Alice'
         >>> scenario_list[1]['age']
         '25'
+
+        >>> url = "https://example.com/data.csv"
+        >>> ## scenario_list_from_url = ScenarioList.from_csv(url)
         """
         from edsl.scenarios.Scenario import Scenario
 
-        observations = []
-        with open(filename, "r") as f:
-            reader = csv.reader(f)
+        def is_url(source):
+            try:
+                result = urllib.parse.urlparse(source)
+                return all([result.scheme, result.netloc])
+            except ValueError:
+                return False
+
+        if isinstance(source, str) and is_url(source):
+            with urllib.request.urlopen(source) as response:
+                csv_content = response.read().decode("utf-8")
+            csv_file = StringIO(csv_content)
+        elif isinstance(source, urllib.parse.ParseResult):
+            with urllib.request.urlopen(source.geturl()) as response:
+                csv_content = response.read().decode("utf-8")
+            csv_file = StringIO(csv_content)
+        else:
+            csv_file = open(source, "r")
+
+        try:
+            reader = csv.reader(csv_file)
             header = next(reader)
-            for row in reader:
-                observations.append(Scenario(dict(zip(header, row))))
+            observations = [Scenario(dict(zip(header, row))) for row in reader]
+        finally:
+            csv_file.close()
+
         return cls(observations)
 
     def _to_dict(self, sort=False) -> dict:
@@ -663,6 +962,16 @@ class ScenarioList(Base, UserList, ScenarioListMixin):
 
         return cls([Scenario.from_dict(s) for s in data["scenarios"]])
 
+    @classmethod
+    def from_nested_dict(cls, data: dict) -> ScenarioList:
+        """Create a `ScenarioList` from a nested dictionary."""
+        from edsl.scenarios.Scenario import Scenario
+
+        s = ScenarioList()
+        for key, value in data.items():
+            s.add_list(key, value)
+        return s
+
     def code(self) -> str:
         ## TODO: Refactor to only use the questions actually in the survey
         """Create the Python code representation of a survey."""

edsl/scenarios/ScenarioListExportMixin.py

@@ -41,3 +41,12 @@ class ScenarioListExportMixin(DatasetExportMixin):
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         decorate_methods_from_mixin(cls, DatasetExportMixin)
+
+    def to_docx(self, filename: str):
+        """Export the ScenarioList to a .docx file."""
+        dataset = self.to_dataset()
+        from edsl.results.DatasetTree import Tree
+
+        tree = Tree(dataset)
+        tree.construct_tree()
+        tree.to_docx(filename)

edsl/study/SnapShot.py

@@ -57,10 +57,17 @@ class SnapShot:
         from edsl.Base import Base
         from edsl.study.Study import Study
 
+        def is_edsl_object(obj):
+            package_name = "edsl"
+            cls = obj.__class__
+            module_name = cls.__module__
+            return module_name.startswith(package_name)
+
         for name, value in namespace.items():
             # TODO check this code logic (if there are other objects with to_dict method that are not from edsl)
             if (
-                hasattr(value, "to_dict")
+                is_edsl_object(value)
+                and hasattr(value, "to_dict")
                 and not inspect.isclass(value)
                 and value.__class__ not in [o.__class__ for o in self.exclude]
             ):

edsl/surveys/RuleCollection.py

@@ -1,6 +1,6 @@
 """A collection of rules for a survey."""
 
-from typing import List, Union, Any
+from typing import List, Union, Any, Optional
 from collections import defaultdict, UserList
 
 from edsl.exceptions import (
@@ -24,7 +24,7 @@ NextQuestion = namedtuple(
 class RuleCollection(UserList):
     """A collection of rules for a particular survey."""
 
-    def __init__(self, num_questions: int = None, rules: List[Rule] = None):
+    def __init__(self, num_questions: Optional[int] = None, rules: List[Rule] = None):
         """Initialize the RuleCollection object.
 
         :param num_questions: The number of questions in the survey.