aboutsummaryrefslogtreecommitdiff
path: root/rpki/csv_utils.py
blob: d354da9a0002463f00cbbb99f9532844eeb49833 (plain) (blame)
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
# $Id$
#
# Copyright (C) 2015--2016  Parsons Government Services ("PARSONS")
# Portions copyright (C) 2013--2014  Dragon Research Labs ("DRL")
# Portions copyright (C) 2009--2012  Internet Systems Consortium ("ISC")
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notices and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND PARSONS, DRL, AND ISC DISCLAIM
# ALL WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS.  IN NO EVENT SHALL
# PARSONS, DRL, OR ISC BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR
# CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
# OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

"""
CSV utilities, moved here from myrpki.py.
"""

import csv
import os

class BadCSVSyntax(Exception):
    """
    Bad CSV syntax.
    """

class csv_reader(object):
    """
    Reader for tab-delimited text that's (slightly) friendlier than the
    stock Python csv module (which isn't intended for direct use by
    humans anyway, and neither was this package originally, but that
    seems to be the way that it has evolved...).

    Columns parameter specifies how many columns users of the reader
    expect to see; lines with fewer columns will be padded with None
    values.

    Original API design for this class courtesy of Warren Kumari, but
    don't blame him if you don't like what I did with his ideas.
    """

    def __init__(self, filename, columns = None, min_columns = None, comment_characters = "#;"):
        assert columns is None or isinstance(columns, int)
        assert min_columns is None or isinstance(min_columns, int)
        if columns is not None and min_columns is None:
            min_columns = columns
        self.columns = columns
        self.min_columns = min_columns
        self.comment_characters = comment_characters
        if isinstance(filename, (str, unicode)):
            # Name of a file to open
            self.filename = filename
            self.file = open(filename, "r")
        else:
            # File-like object, already opened
            self.filename = None
            self.file = filename
            
    def __iter__(self):
        line_number = 0
        for line in self.file:
            line_number += 1
            line = line.strip()
            if not line or line[0] in self.comment_characters:
                continue
            fields = line.split()
            if self.min_columns is not None and len(fields) < self.min_columns:
                raise BadCSVSyntax("%s:%d: Not enough columns in line %r" % (self.filename, line_number, line))
            if self.columns is not None and len(fields) > self.columns:
                raise BadCSVSyntax("%s:%d: Too many  columns in line %r" % (self.filename, line_number, line))
            if self.columns is not None and len(fields) < self.columns:
                fields += tuple(None for i in xrange(self.columns - len(fields)))
            yield fields

    def __enter__(self):
        return self

    def __exit__(self, _type, value, traceback):
        self.file.close()

class csv_writer(object):
    """
    Writer object for tab delimited text.  We just use the stock CSV
    module in excel-tab mode for this.

    If "renmwo" is set (default), the file will be written to
    a temporary name and renamed to the real filename after closing.
    """

    def __init__(self, filename, renmwo = True):
        self.filename = filename
        self.renmwo = "%s.~renmwo%d~" % (filename, os.getpid()) if renmwo else filename
        self.file = open(self.renmwo, "w")
        self.writer = csv.writer(self.file, dialect = csv.get_dialect("excel-tab"))

    def __enter__(self):
        return self

    def __exit__(self, _type, value, traceback):
        self.close()

    def close(self):
        """
        Close this writer.
        """

        if self.file is not None:
            self.file.close()
            self.file = None
            if self.filename != self.renmwo:
                os.rename(self.renmwo, self.filename)

    def __getattr__(self, attr):
        """
        Fake inheritance from whatever object csv.writer deigns to give us.
        """

        return getattr(self.writer, attr)