pulsar.schema.definition

  1#
  2# Licensed to the Apache Software Foundation (ASF) under one
  3# or more contributor license agreements.  See the NOTICE file
  4# distributed with this work for additional information
  5# regarding copyright ownership.  The ASF licenses this file
  6# to you under the Apache License, Version 2.0 (the
  7# "License"); you may not use this file except in compliance
  8# with the License.  You may obtain a copy of the License at
  9#
 10#   http://www.apache.org/licenses/LICENSE-2.0
 11#
 12# Unless required by applicable law or agreed to in writing,
 13# software distributed under the License is distributed on an
 14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 15# KIND, either express or implied.  See the License for the
 16# specific language governing permissions and limitations
 17# under the License.
 18#
 19
 20import copy
 21from abc import abstractmethod
 22from collections import OrderedDict
 23from enum import Enum, EnumMeta
 24from six import with_metaclass
 25
 26
 27def _check_record_or_field(x):
 28    if (type(x) is type and not issubclass(x, Record)) \
 29            and not isinstance(x, Field):
 30        raise Exception('Argument ' + x + ' is not a Record or a Field')
 31
 32
 33class RecordMeta(type):
 34    def __new__(metacls, name, parents, dct):
 35        if name != 'Record':
 36            # Do not apply this logic to the base class itself
 37            dct['_fields'] = RecordMeta._get_fields(dct)
 38            dct['_required'] = False
 39        return type.__new__(metacls, name, parents, dct)
 40
 41    @classmethod
 42    def _get_fields(cls, dct):
 43        # Build a set of valid fields for this record
 44        fields = OrderedDict()
 45        for name, value in dct.items():
 46            if issubclass(type(value), EnumMeta):
 47                value = CustomEnum(value)
 48            elif type(value) == RecordMeta:
 49                # We expect an instance of a record rather than the class itself
 50                value = value()
 51
 52            if isinstance(value, Record) or isinstance(value, Field):
 53                fields[name] = value
 54        return fields
 55
 56
 57class Record(with_metaclass(RecordMeta, object)):
 58
 59    # This field is used to set namespace for Avro Record schema.
 60    _avro_namespace = None
 61
 62    # Generate a schema where fields are sorted alphabetically
 63    _sorted_fields = False
 64
 65    def __init__(self, default=None, required_default=False, required=False, *args, **kwargs):
 66        self._required_default = required_default
 67        self._default = default
 68        self._required = required
 69
 70        for k, value in self._fields.items():
 71            if k in kwargs:
 72                if isinstance(value, Record) and isinstance(kwargs[k], dict):
 73                    # Use dict init Record object
 74                    copied = copy.copy(value)
 75                    copied.__init__(**kwargs[k])
 76                    self.__setattr__(k, copied)
 77                elif isinstance(value, Array) and isinstance(kwargs[k], list) and len(kwargs[k]) > 0 \
 78                        and isinstance(value.array_type, Record) and isinstance(kwargs[k][0], dict):
 79                    arr = []
 80                    for item in kwargs[k]:
 81                        copied = copy.copy(value.array_type)
 82                        copied.__init__(**item)
 83                        arr.append(copied)
 84                    self.__setattr__(k, arr)
 85                elif isinstance(value, Map) and isinstance(kwargs[k], dict) and len(kwargs[k]) > 0 \
 86                    and isinstance(value.value_type, Record) and isinstance(list(kwargs[k].values())[0], dict):
 87                    dic = {}
 88                    for mapKey, mapValue in kwargs[k].items():
 89                        copied = copy.copy(value.value_type)
 90                        copied.__init__(**mapValue)
 91                        dic[mapKey] = copied
 92                    self.__setattr__(k, dic)
 93                else:
 94                    # Value was overridden at constructor
 95                    self.__setattr__(k, kwargs[k])
 96            elif isinstance(value, Record):
 97                # Value is a subrecord
 98                self.__setattr__(k, value)
 99            else:
100                # Set field to default value, without revalidating the default value type
101                super(Record, self).__setattr__(k, value.default())
102
103    @classmethod
104    def schema(cls):
105        return cls.schema_info(set())
106
107    @classmethod
108    def schema_info(cls, defined_names):
109        namespace_prefix = ''
110        if cls._avro_namespace is not None:
111            namespace_prefix = cls._avro_namespace + '.'
112        namespace_name = namespace_prefix + cls.__name__
113
114        if namespace_name in defined_names:
115            return namespace_name
116
117        defined_names.add(namespace_name)
118
119        schema = {
120            'type': 'record',
121            'name': str(cls.__name__)
122        }
123        if cls._avro_namespace is not None:
124            schema['namespace'] = cls._avro_namespace
125        schema['fields'] = []
126
127        def get_filed_default_value(value):
128            if isinstance(value, Enum):
129                return value.name
130            else:
131                return value
132
133        if cls._sorted_fields:
134            fields = sorted(cls._fields.keys())
135        else:
136            fields = cls._fields.keys()
137        for name in fields:
138            field = cls._fields[name]
139            field_type = field.schema_info(defined_names) \
140                if field._required else ['null', field.schema_info(defined_names)]
141            schema['fields'].append({
142                'name': name,
143                'default': get_filed_default_value(field.default()),
144                'type': field_type
145            }) if field.required_default() else schema['fields'].append({
146                'name': name,
147                'type': field_type,
148            })
149
150        return schema
151
152    def __setattr__(self, key, value):
153        if key == '_default':
154            super(Record, self).__setattr__(key, value)
155        elif key == '_required_default':
156            super(Record, self).__setattr__(key, value)
157        elif key == '_required':
158            super(Record, self).__setattr__(key, value)
159        else:
160            if key not in self._fields:
161                raise AttributeError('Cannot set undeclared field ' + key + ' on record')
162
163            # Check that type of value matches the field type
164            field = self._fields[key]
165            value = field.validate_type(key, value)
166            super(Record, self).__setattr__(key, value)
167
168    def __eq__(self, other):
169        for field in self._fields:
170            if self.__getattribute__(field) != other.__getattribute__(field):
171                return False
172        return True
173
174    def __ne__(self, other):
175        return not self.__eq__(other)
176
177    def __str__(self):
178        return str(self.__dict__)
179
180    def type(self):
181        return str(self.__class__.__name__)
182
183    def python_type(self):
184        return self.__class__
185
186    def validate_type(self, name, val):
187        if val is None and not self._required:
188            return self.default()
189
190        if not isinstance(val, self.__class__):
191            raise TypeError("Invalid type '%s' for sub-record field '%s'. Expected: %s" % (
192                type(val), name, self.__class__))
193        return val
194
195    def default(self):
196        if self._default is not None:
197            return self._default
198        else:
199            return None
200
201    def required_default(self):
202        return self._required_default
203
204
205class Field(object):
206    def __init__(self, default=None, required=False, required_default=False):
207        if default is not None:
208            default = self.validate_type('default', default)
209        self._default = default
210        self._required_default = required_default
211        self._required = required
212
213    @abstractmethod
214    def type(self):
215        pass
216
217    @abstractmethod
218    def python_type(self):
219        pass
220
221    def validate_type(self, name, val):
222        if val is None and not self._required:
223            return self.default()
224
225        if type(val) != self.python_type():
226            raise TypeError("Invalid type '%s' for field '%s'. Expected: %s" % (type(val), name, self.python_type()))
227        return val
228
229    def schema(self):
230        # For primitive types, the schema would just be the type itself
231        return self.type()
232
233    def schema_info(self, defined_names):
234        return self.type()
235
236    def default(self):
237        return self._default
238
239    def required_default(self):
240        return self._required_default
241
242
243# All types
244
245
246class Null(Field):
247    def type(self):
248        return 'null'
249
250    def python_type(self):
251        return type(None)
252
253    def validate_type(self, name, val):
254        if val is not None:
255            raise TypeError('Field ' + name + ' is set to be None')
256        return val
257
258
259class Boolean(Field):
260    def type(self):
261        return 'boolean'
262
263    def python_type(self):
264        return bool
265
266    def default(self):
267        if self._default is not None:
268            return self._default
269        else:
270            return False
271
272
273class Integer(Field):
274    def type(self):
275        return 'int'
276
277    def python_type(self):
278        return int
279
280    def default(self):
281        if self._default is not None:
282            return self._default
283        else:
284            return None
285
286
287class Long(Field):
288    def type(self):
289        return 'long'
290
291    def python_type(self):
292        return int
293
294    def default(self):
295        if self._default is not None:
296            return self._default
297        else:
298            return None
299
300
301class Float(Field):
302    def type(self):
303        return 'float'
304
305    def python_type(self):
306        return float
307
308    def default(self):
309        if self._default is not None:
310            return self._default
311        else:
312            return None
313
314
315class Double(Field):
316    def type(self):
317        return 'double'
318
319    def python_type(self):
320        return float
321
322    def default(self):
323        if self._default is not None:
324            return self._default
325        else:
326            return None
327
328
329class Bytes(Field):
330    def type(self):
331        return 'bytes'
332
333    def python_type(self):
334        return bytes
335
336    def default(self):
337        if self._default is not None:
338            return self._default
339        else:
340            return None
341
342
343class String(Field):
344    def type(self):
345        return 'string'
346
347    def python_type(self):
348        return str
349
350    def validate_type(self, name, val):
351        t = type(val)
352
353        if val is None and not self._required:
354            return self.default()
355
356        if not (t is str or t.__name__ == 'unicode'):
357            raise TypeError("Invalid type '%s' for field '%s'. Expected a string" % (t, name))
358        return val
359
360    def default(self):
361        if self._default is not None:
362            return self._default
363        else:
364            return None
365
366# Complex types
367
368
369class CustomEnum(Field):
370    def __init__(self, enum_type, default=None, required=False, required_default=False):
371        if not issubclass(enum_type, Enum):
372            raise Exception(enum_type + " is not a valid Enum type")
373        self.enum_type = enum_type
374        self.values = {}
375        for x in enum_type.__members__.values():
376            self.values[x.value] = x
377        super(CustomEnum, self).__init__(default, required, required_default)
378
379    def type(self):
380        return 'enum'
381
382    def python_type(self):
383        return self.enum_type
384
385    def validate_type(self, name, val):
386        if val is None:
387            return None
388
389        if type(val) is str:
390            # The enum was passed as a string, we need to check it against the possible values
391            if val in self.enum_type.__members__:
392                return self.enum_type.__members__[val]
393            else:
394                raise TypeError(
395                    "Invalid enum value '%s' for field '%s'. Expected: %s" % (val, name, self.enum_type.__members__.keys()))
396        elif type(val) is int:
397            # The enum was passed as an int, we need to check it against the possible values
398            if val in self.values:
399                return self.values[val]
400            else:
401                raise TypeError(
402                    "Invalid enum value '%s' for field '%s'. Expected: %s" % (val, name, self.values.keys()))
403        elif type(val) != self.python_type():
404            raise TypeError("Invalid type '%s' for field '%s'. Expected: %s" % (type(val), name, self.python_type()))
405        else:
406            return val
407
408    def schema(self):
409        return self.schema_info(set())
410
411    def schema_info(self, defined_names):
412        if self.enum_type.__name__ in defined_names:
413            return self.enum_type.__name__
414        defined_names.add(self.enum_type.__name__)
415        return {
416            'type': self.type(),
417            'name': self.enum_type.__name__,
418            'symbols': [x.name for x in self.enum_type]
419        }
420
421    def default(self):
422        if self._default is not None:
423            return self._default
424        else:
425            return None
426
427
428class Array(Field):
429    def __init__(self, array_type, default=None, required=False, required_default=False):
430        _check_record_or_field(array_type)
431        self.array_type = array_type
432        super(Array, self).__init__(default=default, required=required, required_default=required_default)
433
434    def type(self):
435        return 'array'
436
437    def python_type(self):
438        return list
439
440    def validate_type(self, name, val):
441        if val is None:
442            return None
443
444        super(Array, self).validate_type(name, val)
445
446        for x in val:
447            if type(x) != self.array_type.python_type():
448                raise TypeError('Array field ' + name + ' items should all be of type '
449                                + self.array_type.python_type())
450        return val
451
452    def schema(self):
453        return self.schema_info(set())
454
455    def schema_info(self, defined_names):
456        return {
457            'type': self.type(),
458            'items': self.array_type.schema_info(defined_names) if isinstance(self.array_type, (Array, Map, Record))
459                else self.array_type.type()
460        }
461
462    def default(self):
463        if self._default is not None:
464            return self._default
465        else:
466            return None
467
468
469class Map(Field):
470    def __init__(self, value_type, default=None, required=False, required_default=False):
471        _check_record_or_field(value_type)
472        self.value_type = value_type
473        super(Map, self).__init__(default=default, required=required, required_default=required_default)
474
475    def type(self):
476        return 'map'
477
478    def python_type(self):
479        return dict
480
481    def validate_type(self, name, val):
482        if val is None:
483            return None
484
485        super(Map, self).validate_type(name, val)
486
487        for k, v in val.items():
488            if type(k) != str and not is_unicode(k):
489                raise TypeError('Map keys for field ' + name + '  should all be strings')
490            if type(v) != self.value_type.python_type():
491                raise TypeError('Map values for field ' + name + ' should all be of type '
492                                + self.value_type.python_type())
493
494        return val
495
496    def schema(self):
497        return self.schema_info(set())
498
499    def schema_info(self, defined_names):
500        return {
501            'type': self.type(),
502            'values': self.value_type.schema_info(defined_names) if isinstance(self.value_type, (Array, Map, Record))
503                else self.value_type.type()
504        }
505
506    def default(self):
507        if self._default is not None:
508            return self._default
509        else:
510            return None
511
512
513# Python3 has no `unicode` type, so here we use a tricky way to check if the type of `x` is `unicode` in Python2
514# and also make it work well with Python3.
515def is_unicode(x):
516    return 'encode' in dir(x) and type(x.encode()) == str