| 11 | from tqdm import tqdm |
| 12 | |
| 13 | class TSVFile(object): |
| 14 | def __init__(self, tsv_file, total_num): |
| 15 | self.tsv_file = tsv_file |
| 16 | self._fp = None |
| 17 | self.pid = None |
| 18 | self.total_num = total_num |
| 19 | |
| 20 | def __del__(self): |
| 21 | if self._fp: |
| 22 | self._fp.close() |
| 23 | |
| 24 | def __str__(self): |
| 25 | return "TSVFile(tsv_file='{}')".format(self.tsv_file) |
| 26 | |
| 27 | def __repr__(self): |
| 28 | return str(self) |
| 29 | |
| 30 | def num_rows(self): |
| 31 | return self.total_num |
| 32 | |
| 33 | def seek(self, pos): |
| 34 | self._ensure_tsv_opened() |
| 35 | self._fp.seek(pos) |
| 36 | return [s.strip() for s in self._fp.readline().split('\t')] |
| 37 | |
| 38 | def __getitem__(self, pos): |
| 39 | return self.seek(pos) |
| 40 | |
| 41 | def __len__(self): |
| 42 | return self.num_rows() |
| 43 | |
| 44 | |
| 45 | def _ensure_tsv_opened(self): |
| 46 | if self._fp is None: |
| 47 | self._fp = open(self.tsv_file, 'r') |
| 48 | self.pid = os.getpid() |
| 49 | |
| 50 | if self.pid != os.getpid(): |
| 51 | logging.info('re-open {} because the process id changed'.format(self.tsv_file)) |
| 52 | self._fp = open(self.tsv_file, 'r') |
| 53 | self.pid = os.getpid() |
| 54 | |