__init__.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521
  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. ticket_url = None,
  310. valid_from = None,
  311. valid_until = None,
  312. **kwargs
  313. ):
  314. super(Transportation, self).__init__(**kwargs)
  315. if departure_point is not None:
  316. assert type(departure_point) is unicode
  317. self.departure_point = departure_point
  318. if destination_point is not None:
  319. assert type(destination_point) is unicode
  320. self.destination_point = destination_point
  321. if distance is not None:
  322. assert type(distance) is Distance
  323. self.distance = distance
  324. if route_map is not None:
  325. assert type(route_map) is str
  326. self.route_map = route_map
  327. if passenger is not None:
  328. assert type(passenger) is Person
  329. self.passenger = passenger
  330. if valid_from is not None:
  331. assert type(valid_from) is datetime.datetime
  332. assert not valid_from.tzinfo is None
  333. self.valid_from = valid_from
  334. if valid_until is not None:
  335. assert type(valid_until) is datetime.datetime
  336. assert not valid_until.tzinfo is None
  337. self.valid_until = valid_until
  338. if ticket_url is not None:
  339. assert type(ticket_url) is unicode
  340. self.ticket_url = ticket_url
  341. def dict_repr(self):
  342. attr = super(Transportation, self).dict_repr()
  343. attr.update({
  344. 'departure_point': self.departure_point,
  345. 'destination_point': self.destination_point,
  346. 'distance_metres': self.distance.metres if self.distance else None,
  347. 'route_map': self.route_map,
  348. })
  349. return attr
  350. class TaxiRide(Transportation):
  351. yaml_tag = u'!taxi-ride'
  352. def __init__(
  353. self,
  354. arrival_time = None,
  355. departure_time = None,
  356. driver = None,
  357. name = None,
  358. **kwargs
  359. ):
  360. if name is None:
  361. name = u'Taxi Ride'
  362. super(TaxiRide, self).__init__(name = name, **kwargs)
  363. assert type(driver) is unicode
  364. self.driver = driver
  365. assert arrival_time is None or type(arrival_time) is datetime.datetime
  366. self.arrival_time = arrival_time
  367. assert departure_time is None or type(departure_time) is datetime.datetime
  368. self.departure_time = departure_time
  369. def dict_repr(self):
  370. attr = super(TaxiRide, self).dict_repr()
  371. attr.update({
  372. 'arrival_time': self.arrival_time.strftime('%Y-%m-%d %H:%M') if self.arrival_time else None,
  373. 'departure_time': self.departure_time.strftime('%Y-%m-%d %H:%M') if self.departure_time else None,
  374. 'driver': self.driver,
  375. })
  376. return attr
  377. class OrderRegistry(yaml.YAMLObject):
  378. yaml_tag = u'!order-registry'
  379. def __init__(self):
  380. self.registry = {}
  381. def register(self, order):
  382. assert isinstance(order, Order)
  383. if not order.platform in self.registry:
  384. self.registry[order.platform] = {}
  385. self.registry[order.platform][order.order_id] = order
  386. @classmethod
  387. def to_yaml(cls, dumper, self):
  388. return dumper.represent_mapping(cls.yaml_tag, self.registry)
  389. @classmethod
  390. def from_yaml(cls, loader, node):
  391. self = cls()
  392. self.registry = loader.construct_mapping(node)
  393. return self
  394. def __eq__(self, other):
  395. return type(self) == type(other) and vars(self) == vars(other)
  396. def __ne__(self, other):
  397. return not self == other
  398. class Person(_YamlInitConstructor):
  399. yaml_tag = u'!person'
  400. def __init__(self, first_name = None, last_name = None):
  401. self.first_name = first_name
  402. self.last_name = last_name
  403. @property
  404. def first_name(self):
  405. return self._first_name
  406. @first_name.setter
  407. def first_name(self, first_name):
  408. assert first_name is None or type(first_name) is unicode
  409. self._first_name = first_name
  410. @property
  411. def last_name(self):
  412. return self._last_name
  413. @last_name.setter
  414. def last_name(self, last_name):
  415. assert last_name is None or type(last_name) is unicode
  416. self._last_name = last_name
  417. @classmethod
  418. def to_yaml(cls, dumper, person):
  419. return dumper.represent_mapping(cls.yaml_tag, {
  420. 'first_name': person.first_name,
  421. 'last_name': person.last_name,
  422. })
  423. def __eq__(self, other):
  424. return type(self) == type(other) and vars(self) == vars(other)
  425. def __ne__(self, other):
  426. return not self == other
  427. def __repr__(self):
  428. return self.__class__.__name__ + '(%s)' % ', '.join([
  429. '%s=%r' % (k, v) for k, v in vars(self).items()
  430. ])