Graber

Graber

Woskoboy

from weblib.error import DataNotFound

from grab.spider import Spider, Task


import json

import re


from telebot import types


from logging import getLogger

logger = getLogger(__name__)


TRAIN = u'\U0001F68A'

DOWN = u'\U00002B07'

FINISH = u'\U0001F3C1'

TIME_WAY = u'\U000023F3'

STAR = u'\U00002B50'

SEAT = u'\U0001F6B9'



def url_keyboard(url):

  keyboard = types.InlineKeyboardMarkup()

  url_button = types.InlineKeyboardButton(text="Перейти на сайт", url=url)

  keyboard.add(url_button)

  return keyboard



class TutuTrains(Spider):


  def __init__(self, bot, chat_id, url):

    super().__init__()

    self.url = url

    self.bot = bot

    self.chat_id = chat_id


  def task_generator(self):

    yield Task('ya_trains', url=self.url)


  def task_ya_trains(self, grab, task):

    try:

      body = grab.doc.rex_search('\{"componentData.+"searchResultList"\:\[(\{"trains.+),"timeIntervalAlert',

                    flags=re.DOTALL).group(1)+"}"


      for train in json.loads(body)['trains']:

        self.bot.send_message(self.chat_id, get_way(train))


    except DataNotFound:

      self.bot.send_message(self.chat_id, 'Что-то пошло не так. '

                        'Возможно на указанную дату нет прямых рейсов.'

                        ' Попробуй задать другую дату, например, командой '

                        '/d 2017-01-25 или начни сначала командой: /start')



def get_way(train_):


  train = train_['params']['trip']

  run = train_['params']['run']

  travel_h, travel_hh = divmod(train_['params']['trip']['travelTimeSeconds'], 3600)


  try:

    place_price = train_['params']['withSeats']['categories']

    price = get_price(place_price)

  except KeyError:

    price = '\nНет билетов'


  return '{train}{name} {num}.\n{title}-{rate}{star}\n{dep_date} , {dep_time} {down}' \

      '\n{arr_date} , {arr_time} {finish}{price}\n{travel}{travel_h}ч. {travel_m} мин.\n' \

      .format(

        down=DOWN, finish=FINISH, travel=TIME_WAY, star=STAR, train=TRAIN,

        num=train['trainNumber'],

        rate='%s' % run['rating']['value'],

        title=run['rating']['description'][0]['title'],

        dep_date=train['departureDate'],

        dep_time=train['departureTime'][:-3],

        arr_date=train['arrivalDate'],

        arr_time=train['arrivalTime'][:-3],

        price=price,

        travel_h=travel_h,

        travel_m=travel_hh / 60,

        name='Поезд' if not run['name'] else run['name'])



def get_price(elements):


  prices = []

  for cat in elements:

    name = cat['params']['name']

    try:

      price_ = cat['params']['price']['RUB']

      seats_ = cat['params']['seatsCount']

    except KeyError:

      seats_ = '0'

      price_ = '0'


    prices.append('\n{}: {} {} {} руб.'.format(name, seats_, SEAT, price_))


  return ' '.join(prices)



class CitiesKeyboard(Spider):


  def __init__(self, bot, message):

    super().__init__()


    self.bot = bot

    self.text = message.text

    self.chat_id = message.chat.id


  def task_generator(self):


    url = 'http://www.tutu.ru/suggest/railway_simple/?name=%s' % self.text

    yield Task('city_code', url=url, city=self.text)


  def task_city_code(self, grab, task):


    stations = [(_['value'], _['id']) for _ in json.loads(grab.doc.unicode_body())]


    keyboard = create_keyboard(stations)

    self.bot.send_message(self.chat_id, 'Выбери станцию', reply_markup=keyboard)



def create_keyboard(stations):

  keyboard = types.InlineKeyboardMarkup()

  for town, code in stations:

    keyboard.add(types.InlineKeyboardButton(text=town, callback_data=code))

  return keyboard



def log(string):

  logger.warning(string)


Report Page