__init__.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512
  1. # -*- coding: utf-8 -*-
  2. import re
  3. import pytz
  4. import yaml
  5. import yaml.representer
  6. import datetime
  7. yaml.Dumper.add_representer(unicode, yaml.representer.SafeRepresenter.represent_unicode)
  8. yaml.Loader.add_constructor(
  9. u'tag:yaml.org,2002:str',
  10. lambda loader, node: unicode(loader.construct_scalar(node)),
  11. )
  12. def yaml_construct_timestamp(loader, node):
  13. loaded_dt = loader.construct_yaml_timestamp(node)
  14. if type(loaded_dt) is datetime.datetime and loaded_dt.tzinfo is None:
  15. timezone_match = re.search(
  16. ur'\+(?P<h>\d{2}):(?P<m>\d{2})$',
  17. loader.construct_python_unicode(node),
  18. )
  19. if timezone_match:
  20. loaded_dt = loaded_dt.replace(tzinfo = pytz.utc)
  21. return loaded_dt
  22. yaml.Loader.add_constructor(u'tag:yaml.org,2002:timestamp', yaml_construct_timestamp)
  23. class _YamlInitConstructor(yaml.YAMLObject):
  24. @classmethod
  25. def from_yaml(cls, loader, node):
  26. return cls(**loader.construct_mapping(node, deep = True))
  27. # return cls(**{
  28. # k: unicode(v) if isinstance(v, str) else v
  29. # for (k, v) in loader.construct_mapping(node, deep = True).items()
  30. # })
  31. class Figure(_YamlInitConstructor):
  32. yaml_tag = u"!figure"
  33. def __init__(self, value, unit):
  34. self.value = value
  35. self.unit = unit
  36. def get_value(self):
  37. return self._value
  38. def set_value(self, value):
  39. self._value = value
  40. """ use property() instead of decorator to enable overriding """
  41. value = property(get_value, set_value)
  42. def get_unit(self):
  43. return self._unit
  44. def set_unit(self, unit):
  45. assert type(unit) is unicode
  46. self._unit = unit
  47. """ use property() instead of decorator to enable overriding """
  48. unit = property(get_unit, set_unit)
  49. def __eq__(self, other):
  50. return type(self) == type(other) and self.value == other.value and self.unit == other.unit
  51. def __ne__(self, other):
  52. return not (self == other)
  53. @classmethod
  54. def to_yaml(cls, dumper, figure):
  55. return dumper.represent_mapping(
  56. cls.yaml_tag,
  57. {'unit': figure.get_unit(), 'value': figure.get_value()},
  58. )
  59. class ScalarFigure(Figure):
  60. yaml_tag = u'!scalar'
  61. def get_value(self):
  62. return super(ScalarFigure, self).get_value()
  63. def set_value(self, value):
  64. assert type(value) is float
  65. super(ScalarFigure, self).set_value(value)
  66. """ use property() instead of decorator to enable overriding """
  67. value = property(get_value, set_value)
  68. @classmethod
  69. def from_yaml(cls, loader, node):
  70. attr = loader.construct_scalar(node).split(' ')
  71. return cls(
  72. value = float(attr[0]),
  73. unit = attr[1],
  74. )
  75. @classmethod
  76. def to_yaml(cls, dumper, s):
  77. return dumper.represent_scalar(
  78. cls.yaml_tag,
  79. '%s %s' % (repr(s.value), s.unit),
  80. )
  81. class Distance(ScalarFigure):
  82. yaml_tag = u'!distance'
  83. @property
  84. def metres(self):
  85. if self.unit == 'km':
  86. return self.value * 1000
  87. else:
  88. raise Exception()
  89. class Sum(ScalarFigure):
  90. yaml_tag = u'!sum'
  91. def __init__(self, value, currency):
  92. super(Sum, self).__init__(value, currency)
  93. @property
  94. def currency(self):
  95. return self.unit
  96. def get_unit(self):
  97. return super(Sum, self).get_unit()
  98. def set_unit(self, currency):
  99. if currency == u'€':
  100. currency = u'EUR'
  101. if currency == u'US$':
  102. currency = u'USD'
  103. assert type(currency) is unicode
  104. assert currency in [u'EUR', u'USD']
  105. super(Sum, self).set_unit(currency)
  106. """ use property() instead of decorator to enable overriding """
  107. unit = property(get_unit, set_unit)
  108. @classmethod
  109. def from_yaml(cls, loader, node):
  110. attr = loader.construct_scalar(node).split(' ')
  111. return cls(
  112. value = float(attr[0]),
  113. currency = attr[1],
  114. )
  115. class Discount(yaml.YAMLObject):
  116. yaml_tag = u'!discount'
  117. def __init__(
  118. self,
  119. name = None,
  120. amount = None,
  121. ):
  122. assert type(name) is unicode
  123. self.name = name
  124. assert type(amount) is Sum
  125. assert amount.value >= 0
  126. self.amount = amount
  127. def dict_repr(self):
  128. return {
  129. 'name': self.name,
  130. 'value': self.amount.value,
  131. 'value_currency': self.amount.currency,
  132. }
  133. @staticmethod
  134. def from_dict(attr):
  135. return Discount(
  136. name = attr['name'],
  137. amount = Sum(attr['value'], attr['value_currency']),
  138. )
  139. def __eq__(self, other):
  140. return (type(self) == type(other)
  141. and vars(self) == vars(other))
  142. def __ne__(self, other):
  143. return not (self == other)
  144. class Order(_YamlInitConstructor):
  145. yaml_tag = u'!order'
  146. def __init__(self, platform, order_id, order_date,
  147. customer_id = None,
  148. items = None,
  149. discounts = None,
  150. ):
  151. assert type(platform) is unicode
  152. self.platform = platform
  153. if type(order_id) in [int]:
  154. order_id = unicode(order_id)
  155. assert type(order_id) is unicode
  156. self.order_id = order_id
  157. if type(order_date) is datetime.datetime:
  158. order_date = order_date.date()
  159. assert type(order_date) is datetime.date
  160. self.order_date = order_date
  161. assert customer_id is None or type(customer_id) is unicode
  162. self.customer_id = customer_id
  163. if items is None:
  164. items = []
  165. assert type(items) is list
  166. self.items = items
  167. if discounts is None:
  168. discounts = []
  169. assert type(discounts) is list
  170. self.discounts = discounts
  171. def dict_repr(self):
  172. return {k: v for (k, v) in {
  173. 'articles': self.items,
  174. 'customer_id': self.customer_id,
  175. 'discounts': self.discounts,
  176. 'order_date': self.order_date.strftime('%Y-%m-%d'),
  177. 'order_id': self.order_id,
  178. 'platform': self.platform,
  179. }.items() if v is not None}
  180. @staticmethod
  181. def from_dict(attr):
  182. order = Order(
  183. platform = attr['platform'],
  184. order_id = attr['order_id'],
  185. order_date = datetime.datetime.strptime(attr['order_date'], '%Y-%m-%d'),
  186. )
  187. if 'customer_id' in attr:
  188. order.customer_id = attr['customer_id']
  189. for item in attr['articles']:
  190. if type(item) is dict:
  191. item = Item.from_dict(item)
  192. assert isinstance(item, Item)
  193. order.items.append(item)
  194. for discount in attr['discounts']:
  195. if type(discount) is dict:
  196. discount = Discount.from_dict(discount)
  197. assert isinstance(discount, Discount)
  198. order.discounts.append(discount)
  199. return order
  200. def __eq__(self, other):
  201. return (type(self) == type(other)
  202. and vars(self) == vars(other))
  203. def __ne__(self, other):
  204. return not (self == other)
  205. class Item(_YamlInitConstructor):
  206. yaml_tag = u'!item'
  207. def __init__(
  208. self,
  209. name = None,
  210. price_brutto = None,
  211. ):
  212. assert type(name) is unicode
  213. self.name = name
  214. assert type(price_brutto) is Sum
  215. self.price_brutto = price_brutto
  216. def dict_repr(self):
  217. return {
  218. 'name': self.name,
  219. 'price_brutto': self.price_brutto.value,
  220. 'price_brutto_currency': self.price_brutto.currency,
  221. }
  222. @staticmethod
  223. def from_dict(attr):
  224. return Item(
  225. name = attr['name'],
  226. price_brutto = Sum(attr['price_brutto'], attr['price_brutto_currency']),
  227. )
  228. def __eq__(self, other):
  229. return (type(self) == type(other)
  230. and vars(self) == vars(other))
  231. def __ne__(self, other):
  232. return not (self == other)
  233. class Article(Item):
  234. yaml_tag = u'!article'
  235. def __init__(
  236. self,
  237. authors = None,
  238. color = None,
  239. delivery_date = None,
  240. depth = None,
  241. features = None,
  242. height = None,
  243. maximum_load = None,
  244. option = None,
  245. product_id = None,
  246. quantity = None,
  247. reseller = None,
  248. shipper = None,
  249. size = None,
  250. state = None,
  251. width = None,
  252. **kwargs
  253. ):
  254. super(Article, self).__init__(**kwargs)
  255. assert type(quantity) is int
  256. self.quantity = quantity
  257. if authors is not None:
  258. assert type(authors) is list
  259. self.authors = authors
  260. if state is not None:
  261. assert type(state) is unicode
  262. self.state = state
  263. if reseller is not None:
  264. assert type(reseller) is unicode
  265. self.reseller = reseller
  266. if shipper is not None:
  267. assert type(shipper) is unicode
  268. self.shipper = shipper
  269. if product_id is not None:
  270. if type(product_id) in [int]:
  271. product_id = unicode(product_id)
  272. assert type(product_id) is unicode
  273. self.product_id = product_id
  274. if option is not None:
  275. assert type(option) is unicode
  276. self.option = option
  277. if color is not None:
  278. assert type(color) is unicode
  279. self.color = color
  280. if size is not None:
  281. assert type(size) is unicode
  282. self.size = size
  283. if width is not None:
  284. assert type(width) is ScalarFigure
  285. self.width = width
  286. if depth is not None:
  287. assert type(depth) is ScalarFigure
  288. self.depth = depth
  289. if height is not None:
  290. assert type(height) is ScalarFigure
  291. self.height = height
  292. if maximum_load is not None:
  293. assert type(maximum_load) is ScalarFigure
  294. self.maximum_load = maximum_load
  295. if features is not None:
  296. assert type(features) is unicode
  297. self.features = features
  298. assert delivery_date is None or type(delivery_date) is datetime.date
  299. self.delivery_date = delivery_date
  300. class Transportation(Item):
  301. yaml_tag = u'!transportation'
  302. def __init__(
  303. self,
  304. departure_point = None,
  305. destination_point = None,
  306. distance = None,
  307. passenger = None,
  308. route_map = None,
  309. valid_from = None,
  310. **kwargs
  311. ):
  312. super(Transportation, self).__init__(**kwargs)
  313. if departure_point is not None:
  314. assert type(departure_point) is unicode
  315. self.departure_point = departure_point
  316. if destination_point is not None:
  317. assert type(destination_point) is unicode
  318. self.destination_point = destination_point
  319. if distance is not None:
  320. assert type(distance) is Distance
  321. self.distance = distance
  322. if route_map is not None:
  323. assert type(route_map) is str
  324. self.route_map = route_map
  325. if passenger is not None:
  326. assert type(passenger) is Person
  327. self.passenger = passenger
  328. if valid_from is not None:
  329. assert type(valid_from) is datetime.datetime
  330. assert not valid_from.tzinfo is None
  331. self.valid_from = valid_from
  332. def dict_repr(self):
  333. attr = super(Transportation, self).dict_repr()
  334. attr.update({
  335. 'departure_point': self.departure_point,
  336. 'destination_point': self.destination_point,
  337. 'distance_metres': self.distance.metres if self.distance else None,
  338. 'route_map': self.route_map,
  339. })
  340. return attr
  341. class TaxiRide(Transportation):
  342. yaml_tag = u'!taxi-ride'
  343. def __init__(
  344. self,
  345. arrival_time = None,
  346. departure_time = None,
  347. driver = None,
  348. name = None,
  349. **kwargs
  350. ):
  351. if name is None:
  352. name = u'Taxi Ride'
  353. super(TaxiRide, self).__init__(name = name, **kwargs)
  354. assert type(driver) is unicode
  355. self.driver = driver
  356. assert arrival_time is None or type(arrival_time) is datetime.datetime
  357. self.arrival_time = arrival_time
  358. assert departure_time is None or type(departure_time) is datetime.datetime
  359. self.departure_time = departure_time
  360. def dict_repr(self):
  361. attr = super(TaxiRide, self).dict_repr()
  362. attr.update({
  363. 'arrival_time': self.arrival_time.strftime('%Y-%m-%d %H:%M') if self.arrival_time else None,
  364. 'departure_time': self.departure_time.strftime('%Y-%m-%d %H:%M') if self.departure_time else None,
  365. 'driver': self.driver,
  366. })
  367. return attr
  368. class OrderRegistry(yaml.YAMLObject):
  369. yaml_tag = u'!order-registry'
  370. def __init__(self):
  371. self.registry = {}
  372. def register(self, order):
  373. assert isinstance(order, Order)
  374. if not order.platform in self.registry:
  375. self.registry[order.platform] = {}
  376. self.registry[order.platform][order.order_id] = order
  377. @classmethod
  378. def to_yaml(cls, dumper, self):
  379. return dumper.represent_mapping(cls.yaml_tag, self.registry)
  380. @classmethod
  381. def from_yaml(cls, loader, node):
  382. self = cls()
  383. self.registry = loader.construct_mapping(node)
  384. return self
  385. def __eq__(self, other):
  386. return type(self) == type(other) and vars(self) == vars(other)
  387. def __ne__(self, other):
  388. return not self == other
  389. class Person(_YamlInitConstructor):
  390. yaml_tag = u'!person'
  391. def __init__(self, first_name = None, last_name = None):
  392. self.first_name = first_name
  393. self.last_name = last_name
  394. @property
  395. def first_name(self):
  396. return self._first_name
  397. @first_name.setter
  398. def first_name(self, first_name):
  399. assert first_name is None or type(first_name) is unicode
  400. self._first_name = first_name
  401. @property
  402. def last_name(self):
  403. return self._last_name
  404. @last_name.setter
  405. def last_name(self, last_name):
  406. assert last_name is None or type(last_name) is unicode
  407. self._last_name = last_name
  408. @classmethod
  409. def to_yaml(cls, dumper, person):
  410. return dumper.represent_mapping(cls.yaml_tag, {
  411. 'first_name': person.first_name,
  412. 'last_name': person.last_name,
  413. })
  414. def __eq__(self, other):
  415. return type(self) == type(other) and vars(self) == vars(other)
  416. def __ne__(self, other):
  417. return not self == other
  418. def __repr__(self):
  419. return self.__class__.__name__ + '(%s)' % ', '.join([
  420. '%s=%r' % (k, v) for k, v in vars(self).items()
  421. ])