hbertrand commited on
Commit
8e27251
·
unverified ·
1 Parent(s): 5b7d0e6

FIX: drop current (#58)

Browse files

* fix drop current

* fix drop current missing

* fix drop current missing 2

Files changed (1) hide show
  1. buster/documents/pickle.py +5 -1
buster/documents/pickle.py CHANGED
@@ -30,7 +30,11 @@ class DocumentsPickle(DocumentsManager):
30
 
31
  def get_documents(self, source: str) -> pd.DataFrame:
32
  documents = self.documents.copy()
33
- documents = documents[documents.current == 1]
 
 
 
 
34
 
35
  if source is not None and "source" in documents.columns:
36
  documents = documents[documents.source == source]
 
30
 
31
  def get_documents(self, source: str) -> pd.DataFrame:
32
  documents = self.documents.copy()
33
+ if "current" in documents.columns:
34
+ documents = documents[documents.current == 1]
35
+
36
+ # Drop the `current` column
37
+ documents.drop(columns=["current"], inplace=True)
38
 
39
  if source is not None and "source" in documents.columns:
40
  documents = documents[documents.source == source]