| 10 | |
| 11 | |
| 12 | class TSVFile(object): |
| 13 | def __init__(self, tsv_file, silence=True): |
| 14 | self.tsv_file = tsv_file |
| 15 | self.lineidx = op.splitext(tsv_file)[0] + '.lineidx' |
| 16 | |
| 17 | self.label_file = op.splitext(tsv_file)[0] + '.label' |
| 18 | self.label_lineidx = op.splitext(tsv_file)[0] + '.label.lineidx' |
| 19 | |
| 20 | if os.path.exists(self.label_file): |
| 21 | self.split_label = True |
| 22 | else: |
| 23 | self.split_label = False |
| 24 | |
| 25 | self._fp = None |
| 26 | self._lineidx = None |
| 27 | |
| 28 | self._label_fp = None |
| 29 | self._label_lineidx = None |
| 30 | |
| 31 | self.pid = None |
| 32 | self.silence = silence |
| 33 | self._ensure_lineidx_loaded() |
| 34 | |
| 35 | def num_rows(self): |
| 36 | return len(self._lineidx) |
| 37 | |
| 38 | def seek(self, idx): |
| 39 | self._ensure_tsv_opened() |
| 40 | pos = self._lineidx[idx] |
| 41 | self._fp.seek(pos) |
| 42 | tsv_info = [s.strip() for s in self._fp.readline().split('\t')] |
| 43 | |
| 44 | if self.split_label: |
| 45 | label_pos = self._label_lineidx[idx] |
| 46 | self._label_fp.seek(label_pos) |
| 47 | label_info = [s.strip() for s in self._label_fp.readline().split('\t')] |
| 48 | |
| 49 | assert tsv_info[0] == label_info[0] |
| 50 | tsv_info = [tsv_info[0], label_info[-1], tsv_info[-1]] |
| 51 | |
| 52 | return tsv_info |
| 53 | |
| 54 | def close(self): |
| 55 | if self._fp is not None: |
| 56 | self._fp.close() |
| 57 | del self._fp |
| 58 | del self._lineidx |
| 59 | |
| 60 | self._fp = None |
| 61 | self._lineidx = None |
| 62 | |
| 63 | def _ensure_lineidx_loaded(self): |
| 64 | if not op.isfile(self.lineidx) and not op.islink(self.lineidx): |
| 65 | generate_lineidx(self.tsv_file, self.lineidx) |
| 66 | |
| 67 | if self._lineidx is None: |
| 68 | with open(self.lineidx, 'r') as fp: |
| 69 | lines = fp.readlines() |