edsl 0.1.27.dev2__py3-none-any.whl → 0.1.28__py3-none-any.whl

edsl/study/Study.py

@@ -0,0 +1,491 @@
+import os
+import platform
+import socket
+import copy
+import inspect
+import json
+from typing import Optional, List, Dict
+from datetime import datetime
+
+# from edsl.Base import Base
+from edsl import Cache, set_session_cache, unset_session_cache
+from edsl.utilities.utilities import dict_hash
+
+from edsl.study.ObjectEntry import ObjectEntry
+from edsl.study.ProofOfWork import ProofOfWork
+from edsl.study.SnapShot import SnapShot
+
+
+class _StudyFrameMarker:
+    pass
+
+
+class Study:
+    """A study organizes a series of EDSL objects.
+
+    ```python
+    with Study(name = "cool_study") as study:
+        q = QuestionFreeText.example()
+        results = q.run()
+    ```
+
+    The `study` object is a context manager.
+    It lets you group a series of events and objects together.
+
+    >>> with Study(name = "cool_study", verbose = False) as study:
+    ...     from edsl import QuestionFreeText
+    ...     q = QuestionFreeText.example()
+    >>> len(study.objects)
+    1
+    >>> _ = os.system("rm cool_study.json")
+
+
+    It records all the edsl objects that are created during the study.
+    On exit, is saves them to a study file.
+
+    """
+
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        filename: Optional[str] = None,
+        description: Optional[str] = None,
+        objects: Optional[Dict[str, ObjectEntry]] = None,
+        cache: Optional[Cache] = None,
+        coop: bool = False,
+        use_study_cache=True,
+        overwrite_on_change=True,
+        proof_of_work=None,
+        proof_of_work_difficulty: int = None,
+        namespace: Optional[dict] = None,
+        verbose=True,
+    ):
+        """
+        :param name: The name of the study.
+        :param description: A description of the study.
+        :param objects: A dictionary of objects to add to the study.
+        :param cache: A cache object to (potentially) use for the study.
+        :param filename: The path to the study file.
+        :param coop: Whether to push the study to coop.
+        :param use_study_cache: Whether to use the study cache.
+        :param overwrite_on_change: Whether to overwrite the study file if it has changed.
+
+        >>> s = Study()
+        Traceback (most recent call last):
+        ...
+        ValueError: You must provide a name or a filename for the study.
+
+        """
+        self.verbose = verbose
+
+        if name is None and filename is None:
+            raise ValueError("You must provide a name or a filename for the study.")
+
+        if filename is None:
+            self.filename = name
+        else:
+            self.filename = filename
+
+        if (
+            self.filename
+            and os.path.exists(self.filename + ".json")
+            and os.path.getsize(self.filename + ".json") > 0
+        ):
+            if self.verbose:
+                print(f"Using existing study file {self.filename}.json")
+            self._load_from_file()
+        else:
+            self.name = name
+            self.description = description
+            self.objects = objects or {}
+            self.cache = cache or Cache()
+            self.proof_of_work = proof_of_work or ProofOfWork()
+
+        # These always overwrite the saved study
+        self.coop = coop
+        self.use_study_cache = use_study_cache
+        self.overwrite_on_change = overwrite_on_change
+        self.proof_of_work_difficulty = proof_of_work_difficulty
+
+        self.starting_objects = copy.deepcopy(self.objects)
+
+        if namespace:
+            self._namespace = namespace
+        else:
+            self._namespace = None
+
+        self._create_mapping_dicts()
+
+    @property
+    def namespace(self):
+        return self._find_stack()
+
+    def _find_stack(self) -> dict:
+        "Finds the frame with the Study context"
+        # if self.verbose:
+        #    self.explore_stacks()
+        frame = inspect.currentframe()
+        candidate_frames = []
+        while frame:
+            if "Study" in frame.f_globals:
+                candidate_frames.append(frame)
+
+            frame = frame.f_back
+
+        found_variables_dict = {}
+        for frame in candidate_frames:
+            found_variables_dict.update(frame.f_globals)
+            found_variables_dict.update(frame.f_locals)
+
+        return found_variables_dict
+
+    def explore_stacks(self):
+        frame = inspect.currentframe()
+        count = 0
+        d = {}
+        while frame:
+            d[count] = "Study" in frame.f_globals.keys()
+            count += 1
+            from rich import print as rprint
+
+            print("Globals:")
+            rprint(frame.f_globals["__name__"])
+            rprint(frame.f_globals.keys())
+            print("Locals:")
+            rprint(frame.f_locals.keys())
+            print("\n")
+            frame = frame.f_back
+        return d
+
+    def _create_mapping_dicts(self):
+        self._name_to_object = {}
+        self._hash_to_name = {}
+        self._name_to_oe = {}
+        name_counts = {}
+        for hash, obj in self.objects.items():
+            new_name = obj.variable_name
+            if obj.variable_name in name_counts:
+                name_counts[obj.variable_name] += 1
+                new_name = obj.variable_name + "_" + str(name_counts[obj.variable_name])
+            else:
+                name_counts[obj.variable_name] = 1
+            self._name_to_object[new_name] = obj.object
+            self._hash_to_name[hash] = new_name
+
+    def __len__(self):
+        return len(self.objects)
+
+    def __eq__(self, other):
+        return self.objects.keys() == other.objects.keys()
+
+    @property
+    def name_to_object(self):
+        """
+        >>> s = Study.example()
+        >>> s.name_to_object
+        {'q': Question('free_text', question_name = \"""how_are_you\""", question_text = \"""How are you?\""")}
+        """
+        self._create_mapping_dicts()
+        return self._name_to_object
+
+    @property
+    def hash_to_name(self):
+        self._create_mapping_dicts()
+        return self._hash_to_name
+
+    def __getattr__(self, name):
+        return self.name_to_object[name]
+
+    @classmethod
+    def from_file(cls, filename: str):
+        """Load a study from a file."""
+        if filename.endswith(".json"):
+            filename = filename[:-5]
+        return cls(filename=filename)
+
+    def _load_from_file(self):
+        """Load the study from a file.
+
+        >>> import tempfile
+        >>> filename = tempfile.NamedTemporaryFile(delete=False)
+        >>> study = Study(name = "poo", filename = filename.name)
+        >>> study.save()
+        Saving study to ...
+        >>> study2 = Study(filename = filename.name)
+        Using existing study file ...
+        >>> study2.name
+        'poo'
+        """
+        with open(self.filename + ".json", "r") as f:
+            d = json.load(f)
+            d["cache"] = Cache.from_dict(d["cache"])
+            d["proof_of_work"] = ProofOfWork.from_dict(d["proof_of_work"])
+            d["objects"] = {
+                hash: ObjectEntry.from_dict(obj_dict)
+                for hash, obj_dict in d["objects"].items()
+            }
+            self.__dict__.update(d)
+
+    def __enter__(self):
+        """
+        >>> s = Study(name = "temp", use_study_cache = True, verbose = False)
+        >>> _ = s.__enter__()
+        >>> from edsl.config import CONFIG
+        >>> hasattr(CONFIG, "EDSL_SESSION_CACHE")
+        True
+        >>> _ = s.__exit__(None, None, None)
+        >>> len(s.objects)
+        0
+        >>> os.remove("temp.json")
+
+        """
+        if self.verbose:
+            print("Existing objects in study:")
+            self.print()
+        snapshot = SnapShot(self.namespace, exclude=[self])
+        if self.use_study_cache:
+            if self.verbose:
+                print("Using study cache.")
+            set_session_cache(self.cache)
+
+        if snapshot.edsl_objects:
+            raise ValueError(
+                "You have EDSL objects in the global namespace.",
+                "Please remove them before starting a study or put under the 'Study' context manager."
+                "Objects found:",
+                snapshot.edsl_objects,
+            )
+        return self
+
+    def __hash__(self) -> int:
+        return dict_hash(list(self.objects.keys()))
+
+    def study_diff(self):
+        ## Need to also report missing.
+        from edsl.BaseDiff import BaseDiff
+
+        raise NotImplementedError("Need to implement this.")
+
+    def print(self):
+        from rich.console import Console
+        from rich.table import Table
+
+        console = Console()
+        table = Table(title="Study")
+        table.add_column("Original Name")
+        table.add_column("Study Name")
+        table.add_column("Class")
+        table.add_column("Description")
+        table.add_column("Hash")
+        table.add_column("Coop info")
+        table.add_column("Created")
+
+        for obj_hash, obj in self.objects.items():
+            url = (
+                ""
+                if not hasattr(obj, "coop_info") or obj.coop_info is None
+                else obj.coop_info.get("url", "")
+            )
+            table.add_row(
+                obj.variable_name,
+                self.hash_to_name[obj_hash],
+                obj.edsl_class_name,
+                obj.description,
+                obj.hash,
+                url,
+                datetime.fromtimestamp(obj.created_at).strftime("%Y-%m-%d %H:%M:%S"),
+            )
+        # Add cache at the end
+        table.add_row(
+            "N/A - Study Cache",
+            "cache",
+            "Cache",
+            f"Cache of study, entries: {len(self.cache)}",
+            str(hash(self.cache)),
+            "N/A",
+            "N/A",
+        )
+        console.print(table)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        snapshot = SnapShot(namespace=self.namespace, exclude=[self])
+        # print("Frame objects are:", snapshot.namespace.keys())
+        # breakpoint()
+        if self.use_study_cache:
+            unset_session_cache()
+
+        for variable_name, object in snapshot.edsl_objects.items():
+            self._add_edsl_object(object=object, variable_name=variable_name)
+
+        if not self.starting_objects:
+            if self.verbose:
+                print(f"New study saved to {self.filename}.json")
+            self.save()
+
+        if self.starting_objects and list(self.starting_objects.keys()) == list(
+            self.objects.keys()
+        ):
+            if self.verbose:
+                print("Study perfectly replicated.")
+        else:
+            if self.verbose:
+                print("Starting hashes:", self.starting_objects.keys())
+                print("Current hashes:", self.objects.keys())
+            if self.starting_objects:
+                missing = set(self.starting_objects.keys()) - set(self.objects.keys())
+                added = set(self.objects.keys()) - set(self.starting_objects.keys())
+                # breakpoint()
+                print("Study did not perfectly replicate.")
+                for hash in missing:
+                    print(f"Missing object: {self.starting_objects[hash]}")
+                for hash in added:
+                    print(f"Added object: {self.objects[hash]}")
+                if self.overwrite_on_change:
+                    print("Overwriting study file.")
+                    self.save()
+                else:
+                    print(
+                        "Please save the study file with a new name or call study iwth 'overwrite_on_change=True' to overwrite the existing study file."
+                    )
+
+        if self.coop:
+            self.push()
+            if self.overwrite_on_change:
+                self.save()
+            else:
+                raise ValueError(
+                    "If you want to push to coop, you must save the study file with a new name or call study iwth 'overwrite_on_change=True' to overwrite the existing study file."
+                )
+
+        if self.verbose:
+            print("Objects in study now:")
+            self.print()
+
+        if self.proof_of_work_difficulty:
+            print("Adding proof of work to study...")
+            from edsl.study.ProofOfWork import ProofOfWork
+
+            # TODO: Need to check if hashes are the same.
+            if not self.proof_of_work.input_data:
+                self.proof_of_work.add_input_data(str(self.__hash__()))
+            self.proof_of_work.add_proof(self.proof_of_work_difficulty)
+            print(
+                "Proof of work added to study with difficulty ",
+                self.proof_of_work_difficulty,
+            )
+            print(self.proof_of_work)
+            self.save()
+
+    def to_dict(self):
+        return {
+            "name": self.name,
+            "description": self.description,
+            "objects": {hash: obj.to_dict() for hash, obj in self.objects.items()},
+            "filename": self.filename,
+            "cache": self.cache.to_dict(),
+            "use_study_cache": self.use_study_cache,
+            "overwrite_on_change": self.overwrite_on_change,
+            "proof_of_work": self.proof_of_work.to_dict(),
+        }
+
+    def versions(self):
+        """Return a dictionary of objects grouped by variable name."""
+        d = {}
+        for _, obj_entry in self.objects.items():
+            if obj_entry.variable_name not in d:
+                d[obj_entry.variable_name] = [obj_entry]
+            else:
+                d[obj_entry.variable_name].append(obj_entry)
+
+        return d
+
+    def diff(self, variable_name: str, index1: int, index2: int):
+        """Return the difference between the versions of an object."""
+        versions = self.versions()[variable_name]
+        diff = versions[index2].object - versions[index1].object
+        return diff
+
+    @classmethod
+    def example(cls, verbose=False):
+        import tempfile
+
+        study_file = tempfile.NamedTemporaryFile()
+        with cls(filename=study_file.name, verbose=verbose) as study:
+            from edsl import QuestionFreeText
+
+            q = QuestionFreeText.example()
+        return study
+
+    @classmethod
+    def from_dict(cls, d):
+        d["cache"] = Cache.from_dict(d["cache"])
+        d["objects"] = {
+            str(object_hash): ObjectEntry.from_dict(obj_dict)
+            for object_hash, obj_dict in d["objects"].items()
+        }
+        d["proof_of_work"] = ProofOfWork.from_dict(d["proof_of_work"])
+        return cls(**d)
+
+    def save(self):
+        if self.verbose:
+            print(f"Saving study to {self.filename}.json")
+        with open(self.filename + ".json", "w") as f:
+            json.dump(self.to_dict(), f, indent=4)
+
+    def _get_system_info(self):
+        return {
+            "platform": platform.platform(),
+            "system": platform.system(),
+            "processor": platform.processor(),
+            "hostname": socket.gethostname(),
+        }
+
+    @staticmethod
+    def _get_description(object):
+        text = ""
+        if hasattr(object, "__len__"):
+            text += f"Num. entries: {len(object)}"
+        if hasattr(object, "question_name"):
+            text += f"Question name: {object.question_name}"
+        return text
+
+    def _add_edsl_object(self, object, variable_name, description=None) -> None:
+        """
+        >>> s = Study.example()
+        >>> from edsl import QuestionLinearScale
+        >>> s._add_edsl_object(QuestionLinearScale.example(), 'q')
+        >>> len(s)
+        2
+        """
+        if description is None:
+            description = self._get_description(object)
+        oe = ObjectEntry(
+            variable_name=variable_name, object=object, description=description
+        )
+        if oe.hash in self.objects:
+            return
+        else:
+            self.objects[oe.hash] = oe
+
+    def push(self, refresh=False) -> None:
+        """Push the objects to coop."""
+        for obj_entry in self.objects.values():
+            obj_entry.push(refresh=refresh)
+
+    def __repr__(self):
+        return f"""Study(name = {self.name}, description = {self.description}, objects = {self.objects}, cache = {self.cache}, filename = {self.filename}, coop = {self.coop}, use_study_cache = {self.use_study_cache}, overwrite_on_change = {self.overwrite_on_change})"""
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(optionflags=doctest.ELLIPSIS)
+
+    # with Study(name = "cool_study") as study:
+    #      from edsl import QuestionFreeText
+    #      q = QuestionFreeText.example()
+
+    # assert len(study.objects) == 1
+
+    # print(study.versions())
+    # {'q': [ObjectEntry(variable_name='q', object=Question('free_text', question_name = """how_are_you""", question_text = """How are you?"""), description='Question name: how_are_you', coop_info=None, created_at=1720276402.561273, edsl_class_name='QuestionFreeText')]}

edsl/study/__init__.py

@@ -0,0 +1,2 @@
+from edsl.study.ObjectEntry import ObjectEntry
+from edsl.study.ProofOfWork import ProofOfWork