From 73a0e14df017dcb0d8ba5c2ed2cfa2d17d581eab Mon Sep 17 00:00:00 2001
From: "J. Sebastian Paez" <jspaezp@gmail.com>
Date: Fri, 6 Dec 2024 15:35:59 -0600
Subject: [PATCH] wip,fix: added deleter to on psm dataset

---
 mokapot/brew.py    | 12 +++++++-----
 mokapot/dataset.py |  6 +++++-
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/mokapot/brew.py b/mokapot/brew.py
index 62ccce9..0526331 100644
--- a/mokapot/brew.py
+++ b/mokapot/brew.py
@@ -5,7 +5,7 @@
 import copy
 import logging
 from operator import itemgetter
-from typing import Iterable
+from typing import Iterable, Generator
 
 import numpy as np
 import pandas as pd
@@ -304,7 +304,9 @@ def brew(
 
 
 # Utility Functions -----------------------------------------------------------
-def make_train_sets(test_idx, subset_max_train, data_size, rng):
+def make_train_sets(
+    test_idx, subset_max_train, data_size, rng
+) -> Generator[list[list[int]], None, None]:
     """
     Parameters
     ----------
@@ -317,8 +319,8 @@ def make_train_sets(test_idx, subset_max_train, data_size, rng):
 
     Yields
     ------
-    PsmDataset
-        The training set.
+    list of list of int
+        The training set. Each element is a list of ints.
     """
     subset_max_train_per_file = []
     if subset_max_train is not None:
@@ -356,7 +358,7 @@ def make_train_sets(test_idx, subset_max_train, data_size, rng):
                 if current_subset_max_train < train_idx_size:
                     train_idx[i] = rng.choice(
                         train_idx[i], current_subset_max_train, replace=False
-                    )
+                    ).tolist()
         yield train_idx
 
 
diff --git a/mokapot/dataset.py b/mokapot/dataset.py
index 5dc46c7..48e6fb9 100644
--- a/mokapot/dataset.py
+++ b/mokapot/dataset.py
@@ -659,12 +659,16 @@ def target_column(self) -> str:
 
     @property
     def feature_columns(self) -> list[str]:
-        return self._feature_columns
+        return list(self._feature_columns)
 
     @property
     def spectra_dataframe(self) -> pd.DataFrame:
         return self._spectra_dataframe
 
+    @spectra_dataframe.deleter
+    def spectra_dataframe(self):
+        del self._spectra_dataframe
+
     def get_column_names(self) -> list[str]:
         columns = self.reader.get_column_names()
         return columns