openrvdas.logger.writers.text_file_writer

  1#!/usr/bin/env python3
  2
  3import os.path
  4import sys
  5import datetime
  6from typing import Union
  7
  8from logger.utils.das_record import DASRecord  # noqa E402
  9from logger.writers.writer import Writer  # noqa: E402
 10
 11
 12class TextFileWriter(Writer):
 13    """Write to the specified file. If filename is empty, write to stdout."""
 14
 15    def __init__(self, filename=None, flush=True, truncate=False,
 16                 split_by_date=False, create_path=True, header=None,
 17                 header_file=None, **kwargs):
 18        """Write text records to a file. If no filename is specified, write to
 19        stdout.
 20        ```
 21        filename     Name of file to write to. If None, write to stdout
 22
 23        flush        If True (default), flush after every write() call
 24
 25        truncate     Truncate file before beginning to write
 26
 27        split_by_date Create a separate text file for every day, appending
 28                     a -YYYY-MM-DD string to the specified filename.
 29
 30        create_path  Create directory path to file if it doesn't exist
 31
 32        header       Add the specified header string to each file.
 33
 34        header_file  Add the content of the specified file to each file.
 35      ```
 36        """
 37        # Initialize type checking
 38        super().__init__(**kwargs)  # processes 'quiet' and type hints
 39
 40        self.filename = filename
 41        self.flush = flush
 42        self.truncate = truncate
 43        self.split_by_date = split_by_date
 44        self.header = None
 45
 46        if split_by_date and not filename:
 47            raise ValueError('TextFileWriter: filename must be specified if '
 48                             'split_by_date is True.')
 49
 50        if header is not None and header_file is not None:
 51            raise ValueError('FileWriter: cannot specify the header and '
 52                             'header_file arguments.')
 53
 54        if header is not None:
 55            if isinstance(header, str):
 56                self.header = header + '\n'
 57            else:
 58                raise ValueError('FileWriter: Unable to add header to data '
 59                                 'file. header argument must be a string: %s',
 60                                 header)
 61
 62        if header_file is not None:
 63            try:
 64                with open(header_file, 'r') as file:
 65                    self.header = file.read()
 66            except:  # noqa E722 - though should really fix this
 67                raise ValueError('FileWriter: Unable to add header to data '
 68                                 'file. header_file argument must be a valid '
 69                                 'filepath: %s', header_file)
 70
 71        # If we're splitting by date, keep track of current file date
 72        # here.
 73        self.file_date = None
 74        self.file = None
 75
 76        # If directory doesn't exist, try to create it
 77        if filename and create_path:
 78            file_dir = os.path.dirname(filename)
 79            if file_dir:
 80                os.makedirs(file_dir, exist_ok=True)
 81
 82        # Figure out what file we ought be writing to and open it.
 83        self._set_file()
 84
 85    ############################
 86    def _today(self):
 87        """Return a tuple for (year, month, day). Broken out into a separate
 88        function to facilitate testing."""
 89        now = datetime.datetime.utcnow()
 90        return (now.year, now.month, now.day)
 91
 92    ############################
 93    def _set_file(self):
 94        """Make sure the right file is open. If we're splitting by date and
 95        the date has rolled over, close the old file and open a new
 96        one. This all feels overly convoluted, but is necessary to keep
 97        checking if we're splitting by date.
 98        """
 99
100        # If they haven't given us a filename, we'll write to stdout
101        if self.filename is None:
102            self.file = sys.stdout
103
104            if self.header is not None:
105                self.file.write(self.header)
106
107            return
108
109        # If here, we have a filename. Check if we're splitting by date;
110        # if so, see if it's time to close out our current file an start a
111        # new one.
112        if self.split_by_date:
113            today = self._today()
114            if self.file_date != today:
115                self.file_date = today
116                if self.file:
117                    self.file.close()
118                    self.file = None
119
120        # If we do have a file open, return.
121        if self.file:
122            return
123
124        # If here, we don't have a file open. This may be because it's our
125        # first time writing, or because we're splitting by dates and
126        # we've rolled over to a new date.
127        if self.split_by_date:
128            today_str = '%04d-%02d-%02d' % self.file_date
129            filename = '%s-%s' % (self.filename, today_str)
130        else:
131            filename = self.filename
132
133        # Open and set the file
134        mode = 'w' if self.truncate else 'a'
135        self.file = open(filename, mode)
136
137        # Add header record to file if a header was specified.
138        if self.header is not None:
139            self.file.write(self.header)
140
141    ############################
142    def write(self, record: Union[str, DASRecord]):
143        """ Write out record, appending a newline at end."""
144
145        if not self.can_process_record(record):  # inherited from BaseModule()
146            self.digest_record(record)           # inherited from BaseModule()
147            return
148
149        if isinstance(record, DASRecord):
150            record = record.as_json()
151
152        # If we're splitting by date, make sure that we're still writing
153        # to the right file.
154        if self.split_by_date:
155            self._set_file()
156
157        # Write the record and flush if requested
158        self.file.write(str(record) + '\n')
159        if self.flush:
160            self.file.flush()