slimmed down dataset.py to only what's

necessary. Updated the structure to the new one

src/utils/dataset.py

@@ -5,13 +5,6 @@ import json, uuid
 
 from utils.errors import InvalidJsonFormatError
 
-# fmt: off
-@dataclass
-class FileData:
-    is_code_related: bool
-    coverage: Dict[str, float] # jacoco-report -> coverage
-    content_before_pr: str = ""
-    content_after_pr: str = ""
 
 @dataclass
 class Comment:
@@ -21,16 +14,18 @@ class Comment:
     to: int
     paraphrases: List[str] = field(default_factory=list)
 
+
 @dataclass
 class Selection:
     comment_suggests_change: bool
     diff_after_address_change: Optional[bool]
-    is_code_related: bool
+
 
 class ArchiveState(Enum):
     BASE = "base"
     MERGED = "merged"
 
+
 @dataclass
 class Metadata:
     id: str
@@ -39,17 +34,20 @@ class Metadata:
     pr_title: str
     pr_body: str
     merge_commit_sha: str   # to checkout for the tests
-    successful: bool = True
+    is_covered: Optional[bool] = None
+    is_code_related: Optional[bool] = None
+    successful: Optional[bool] = None
     build_system: str = ""
     reason_for_failure: str = ""
     last_cmd_error_msg: str = ""
     selection: Optional[Selection] = None
 
-    def archive_name(self, state: ArchiveState, only_id:bool=False):
+    def archive_name(self, state: ArchiveState, only_id: bool = False):
         if only_id:
             return f"{self.id}_{state.value}.tar.gz"
         return f"{self.repo.replace('/', '_')}_{self.pr_number}_{state.value}.tar.gz"
 
+
 @dataclass
 class CommentGenSubmission:
     path: str
@@ -63,7 +61,9 @@ class CommentGenSubmission:
             raise InvalidJsonFormatError("Submitted json doesn't contain an object")
         if not all(k in data and isinstance(data[k], str) for k in ["path", "body"]):
             raise InvalidJsonFormatError("Submitted json doesn't contain the required fields")
-        if not all(k in data and isinstance(data[k], (int, type(None))) for k in ["line_from", "line_to"]):
+        if not all(
+            k in data and isinstance(data[k], (int, type(None))) for k in ["line_from", "line_to"]
+        ):
             raise InvalidJsonFormatError("Submitted json doesn't contain the required fields")
 
         return cls(
@@ -73,52 +73,19 @@ class CommentGenSubmission:
             body=data["body"],
         )
 
+
 @dataclass
 class DatasetEntry:
     metadata: Metadata
-    files: Dict[str, FileData]   # filename -> file data, files before the PR (before the first PR commits)
-    diffs_before: Dict[str, str]   # filename -> diff, diffs between the opening of the PR and the comment
-    comments: List[Comment]
-    diffs_after: Dict[str, str]   # filename -> diff, changes after the comment
-
-
-@dataclass
-class CommentGenEntry:
-    id: str
-    files: Dict[str, str]   # filename -> file content
-    diffs: Dict[str, str]   # filename -> diff, diffs between the opening of the PR and the comment
-
-    @staticmethod
-    def from_entry(entry: DatasetEntry) -> "CommentGenEntry":
-        return CommentGenEntry(
-            id=entry.metadata.id,
-            files={fname: fdata.content_before_pr for fname, fdata in entry.files.items()},
-            diffs=entry.diffs_before,
-        )
-
-
-@dataclass
-class CodeRefinementEntry:
-    id: str
-    files: Dict[str, str]   # filename -> file content
-    diffs: Dict[str, str]   # filename -> diff, diffs between the opening of the PR and the comment
     comments: List[Comment]
 
-    @staticmethod
-    def from_entry(entry: DatasetEntry) -> "CodeRefinementEntry":
-        return CodeRefinementEntry(
-            id=entry.metadata.id,
-            files={fname: fdata.content_before_pr for fname, fdata in entry.files.items()},
-            diffs=entry.diffs_before,
-            comments=entry.comments,
-        )
 
 class OutputType(Enum):
     FULL = "full"
     CODE_REFINEMENT = "code_refinement"
     COMMENT_GEN = "comment_gen"
 
-# fmt: on
+
 @dataclass
 class Dataset:
     entries: List[DatasetEntry] = field(default_factory=list)
@@ -126,59 +93,6 @@ class Dataset:
     def __len__(self) -> int:
         return sum(1 for entry in self.entries if entry.metadata.successful)
 
-    def to_json(
-        self,
-        filename: str,
-        type_: OutputType = OutputType.FULL,
-        remove_non_suggesting: bool = False,
-    ) -> None:
-        """Serialize the dataset to a JSON file"""
-
-        entries_to_dump = self.entries
-
-        if type_ == OutputType.COMMENT_GEN:
-            entries_to_dump = [
-                entry
-                for entry in self.entries
-                if entry.metadata.selection and entry.metadata.selection.comment_suggests_change
-            ]
-        elif type_ == OutputType.CODE_REFINEMENT:
-            entries_to_dump = [
-                entry
-                for entry in self.entries
-                if entry.metadata.selection
-                and entry.metadata.selection.diff_after_address_change
-                and entry.metadata.selection.is_code_related
-            ]
-        elif type_ == OutputType.FULL and remove_non_suggesting:
-            entries_to_dump = [
-                entry
-                for entry in self.entries
-                if entry.metadata.selection and entry.metadata.selection.comment_suggests_change
-            ]
-
-        to_dump = Dataset(entries=entries_to_dump)
-        # print(f"{len(entries_to_dump)} entries...", end=" ", flush=True)
-
-        def transform_entry(entry: Union[DatasetEntry, Dataset, Any]) -> Union[dict, list]:
-            if not isinstance(entry, (DatasetEntry, Dataset)):
-                return entry.__dict__
-
-            if type_ == OutputType.FULL:
-                return entry.__dict__
-
-            if isinstance(entry, Dataset):
-                return entry.entries
-
-            if type_ == OutputType.COMMENT_GEN:
-                return CommentGenEntry.from_entry(entry).__dict__
-
-            if type_ == OutputType.CODE_REFINEMENT:
-                return CodeRefinementEntry.from_entry(entry).__dict__
-
-        with open(filename, "w", encoding="utf-8") as f:
-            json.dump(to_dump, f, default=transform_entry, indent=4)
-
     @staticmethod
     def from_json(filename: str, keep_still_in_progress: bool = False) -> "Dataset":
         with open(filename, "r", encoding="utf-8") as f:
@@ -202,16 +116,11 @@ class Dataset:
             ):
                 continue
 
-            files = {fname: FileData(**fdata) for fname, fdata in entry_data["files"].items()}
-
             comments = [Comment(**comment) for comment in entry_data["comments"]]
 
             entry = DatasetEntry(
                 metadata=metadata,
-                files=files,
-                diffs_before=entry_data["diffs_before"],
                 comments=comments,
-                diffs_after=entry_data["diffs_after"],
             )
             entries.append(entry)