xref: /freebsd/contrib/lib9p/pytest/pfod.py (revision c9ccf3a32da427475985b85d7df023ccfb138c27)
1#! /usr/bin/env python
2
3from __future__ import print_function
4
5__all__ = ['pfod', 'OrderedDict']
6
7### shameless stealing from namedtuple here
8
9"""
10pfod - prefilled OrderedDict
11
12This is basically a hybrid of a class and an OrderedDict,
13or, sort of a data-only class.  When an instance of the
14class is created, all its fields are set to None if not
15initialized.
16
17Because it is an OrderedDict you can add extra fields to an
18instance, and they will be in inst.keys().  Because it
19behaves in a class-like way, if the keys are 'foo' and 'bar'
20you can write print(inst.foo) or inst.bar = 3.  Setting an
21attribute that does not currently exist causes a new key
22to be added to the instance.
23"""
24
25import sys as _sys
26from keyword import iskeyword as _iskeyword
27from collections import OrderedDict
28from collections import deque as _deque
29
30_class_template = '''\
31class {typename}(OrderedDict):
32    '{typename}({arg_list})'
33    __slots__ = ()
34
35    _fields = {field_names!r}
36
37    def __init__(self, *args, **kwargs):
38        'Create new instance of {typename}()'
39        super({typename}, self).__init__()
40        args = _deque(args)
41        for field in self._fields:
42            if field in kwargs:
43                self[field] = kwargs.pop(field)
44            elif len(args) > 0:
45                self[field] = args.popleft()
46            else:
47                self[field] = None
48        if len(kwargs):
49            raise TypeError('unexpected kwargs %s' % kwargs.keys())
50        if len(args):
51            raise TypeError('unconsumed args %r' % tuple(args))
52
53    def _copy(self):
54        'copy to new instance'
55        new = {typename}()
56        new.update(self)
57        return new
58
59    def __getattr__(self, attr):
60        if attr in self:
61            return self[attr]
62        raise AttributeError('%r object has no attribute %r' %
63            (self.__class__.__name__, attr))
64
65    def __setattr__(self, attr, val):
66        if attr.startswith('_OrderedDict_'):
67            super({typename}, self).__setattr__(attr, val)
68        else:
69            self[attr] = val
70
71    def __repr__(self):
72        'Return a nicely formatted representation string'
73        return '{typename}({repr_fmt})'.format(**self)
74'''
75
76_repr_template = '{name}={{{name}!r}}'
77
78# Workaround for py2k exec-as-statement, vs py3k exec-as-function.
79# Since the syntax differs, we have to exec the definition of _exec!
80if _sys.version_info[0] < 3:
81    # py2k: need a real function.  (There is a way to deal with
82    # this without a function if the py2k is new enough, but this
83    # works in more cases.)
84    exec("""def _exec(string, gdict, ldict):
85        "Python 2: exec string in gdict, ldict"
86        exec string in gdict, ldict""")
87else:
88    # py3k: just make an alias for builtin function exec
89    exec("_exec = exec")
90
91def pfod(typename, field_names, verbose=False, rename=False):
92    """
93    Return a new subclass of OrderedDict with named fields.
94
95    Fields are accessible by name.  Note that this means
96    that to copy a PFOD you must use _copy() - field names
97    may not start with '_' unless they are all numeric.
98
99    When creating an instance of the new class, fields
100    that are not initialized are set to None.
101
102    >>> Point = pfod('Point', ['x', 'y'])
103    >>> Point.__doc__                   # docstring for the new class
104    'Point(x, y)'
105    >>> p = Point(11, y=22)             # instantiate with positional args or keywords
106    >>> p
107    Point(x=11, y=22)
108    >>> p['x'] + p['y']                 # indexable
109    33
110    >>> p.x + p.y                       # fields also accessable by name
111    33
112    >>> p._copy()
113    Point(x=11, y=22)
114    >>> p2 = Point()
115    >>> p2.extra = 2
116    >>> p2
117    Point(x=None, y=None)
118    >>> p2.extra
119    2
120    >>> p2['extra']
121    2
122    """
123
124    # Validate the field names.  At the user's option, either generate an error
125    if _sys.version_info[0] >= 3:
126        string_type = str
127    else:
128        string_type = basestring
129    # message or automatically replace the field name with a valid name.
130    if isinstance(field_names, string_type):
131        field_names = field_names.replace(',', ' ').split()
132    field_names = list(map(str, field_names))
133    typename = str(typename)
134    if rename:
135        seen = set()
136        for index, name in enumerate(field_names):
137            if (not all(c.isalnum() or c=='_' for c in name)
138                or _iskeyword(name)
139                or not name
140                or name[0].isdigit()
141                or name.startswith('_')
142                or name in seen):
143                field_names[index] = '_%d' % index
144            seen.add(name)
145    for name in [typename] + field_names:
146        if type(name) != str:
147            raise TypeError('Type names and field names must be strings')
148        if not all(c.isalnum() or c=='_' for c in name):
149            raise ValueError('Type names and field names can only contain '
150                             'alphanumeric characters and underscores: %r' % name)
151        if _iskeyword(name):
152            raise ValueError('Type names and field names cannot be a '
153                             'keyword: %r' % name)
154        if name[0].isdigit():
155            raise ValueError('Type names and field names cannot start with '
156                             'a number: %r' % name)
157    seen = set()
158    for name in field_names:
159        if name.startswith('_OrderedDict_'):
160            raise ValueError('Field names cannot start with _OrderedDict_: '
161                             '%r' % name)
162        if name.startswith('_') and not rename:
163            raise ValueError('Field names cannot start with an underscore: '
164                             '%r' % name)
165        if name in seen:
166            raise ValueError('Encountered duplicate field name: %r' % name)
167        seen.add(name)
168
169    # Fill-in the class template
170    class_definition = _class_template.format(
171        typename = typename,
172        field_names = tuple(field_names),
173        arg_list = repr(tuple(field_names)).replace("'", "")[1:-1],
174        repr_fmt = ', '.join(_repr_template.format(name=name)
175                             for name in field_names),
176    )
177    if verbose:
178        print(class_definition,
179            file=verbose if isinstance(verbose, file) else _sys.stdout)
180
181    # Execute the template string in a temporary namespace and support
182    # tracing utilities by setting a value for frame.f_globals['__name__']
183    namespace = dict(__name__='PFOD%s' % typename,
184                     OrderedDict=OrderedDict, _deque=_deque)
185    try:
186        _exec(class_definition, namespace, namespace)
187    except SyntaxError as e:
188        raise SyntaxError(e.message + ':\n' + class_definition)
189    result = namespace[typename]
190
191    # For pickling to work, the __module__ variable needs to be set to the frame
192    # where the named tuple is created.  Bypass this step in environments where
193    # sys._getframe is not defined (Jython for example) or sys._getframe is not
194    # defined for arguments greater than 0 (IronPython).
195    try:
196        result.__module__ = _sys._getframe(1).f_globals.get('__name__', '__main__')
197    except (AttributeError, ValueError):
198        pass
199
200    return result
201
202if __name__ == '__main__':
203    import doctest
204    doctest.testmod()
205