SSTableConvertMod/MPCORBFileTable.py at main · lsst-dm/SSTableConvertMod · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
from __future__ import annotations

__all__ = ("MPCORBFileTable",)

import csv
from glob import glob
from itertools import islice
from typing import Dict, Optional, Iterable, Generator

from .base import FileTableBuilder, FileTable, Indexer
from .schemas import MPCORB
from .customTypes import ColumnName


class MPCORBBuilder(FileTableBuilder):
    input_schema = ("S3MID", "FORMAT", "q", "e", "i", "Omega",
                    "argperi", "t_p", "H", "t_0", "INDEX",
                    "N_PAR", "MOID", "COMPCODE")

    def __init__(self, parent: FileTable, input_fileglob: str,
                 output_filename: str,
                 skip_rows: int,
                 stop_after: Optional[int] = None,
                 columns: Optional[Iterable[ColumnName]] = None):
        self.parent = parent
        self.output_filename = output_filename
        self.skip_rows = 0
        self.stop_after = None
        self.columns = columns
        self.input_fileglob = input_fileglob
        self._mpc_skip_start = skip_rows
        self._mpc_stop_after = stop_after
        self.do_index = True

        self.index_pos = {self.parent.schema.field_pos[column]: column
                          for column in
                          self.parent.index_columns}

    def _get_input_rows(self) -> Generator:
        if self._mpc_stop_after is not None:
            stop = self._mpc_stop_after + self._mpc_stop_after
        else:
            stop = None  # type: ignore
        for path in glob(self.input_fileglob):
            with open(path, 'rb') as in_file:
                yield from islice(in_file.readlines(), self._mpc_skip_start,
                                  stop)

    def _intrepret_row(self, interp_row: str) -> Dict:
        return {k: v for k, v in zip(self.input_schema,
                                     interp_row.split())}

    def run(self):
        with open(self.output_filename, 'w+', newline='') as out_file:
            indexer = Indexer(self.do_index,
                              self.output_filename+".sidecar",
                              self.parent.index_columns)
            writer = csv.writer(out_file, quoting=csv.QUOTE_NONE,
                                lineterminator="\n")
            writer.writerow(self.parent.schema.fields.keys())
            rows_generator = self._get_input_rows()
            rows = self._make_rows(rows_generator, self.columns,
                                   self.skip_rows,
                                   self.stop_after)
            writer.writerows(indexer.insert(
                (b, i in self.index_pos) for i, b in enumerate(row_gen))
                for row_gen in rows)


class MPCORBFileTable(FileTable):
    schema = MPCORB
    index_columns = tuple(ColumnName(x) for x in ("ssObjectId", "mpcH"))
    builder = MPCORBBuilder