twiins/twiins.py

from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
from tabulate import tabulate
import time
import typer
import click_spinner

app = typer.Typer(help="Minimal Twitch information gatherer")
SEP = '==&&=='


def func_return(result, should_print, pretty=False):
    if should_print:
        if pretty:
            table = []
            for row in result.split('\n'):
                table.append(row.split(SEP))
            print(tabulate(table, tablefmt="plain"))
        else:
            print(result)
    else:
        return result


def get_driver():
    with click_spinner.spinner():
        typer.echo("Providing driver...")
        options = webdriver.FirefoxOptions()
        options.add_argument('--headless')
        options.set_preference('permissions.default.image', 2)
        options.set_preference(
            'dom.ipc.plugins.enabled.libflashplayer.so', 'false')
        driver = webdriver.Firefox(options=options)
        return driver


@app.command()
def check_live(
    channel: str,
    should_print: bool = typer.Option(False, '--print', '-p')
):
    """
    Checks if CHANNEL is livestreaming.
    SHOULDPRINT determines if result is returned to stdout or function caller.
    """
    sel = 'a[data-a-target="watch-mode-to-home"]'
    URL = f'https://www.twitch.tv/{channel}'

    driver = get_driver()
    driver.get(URL)

    is_live = None

    try:
        with click_spinner.spinner():
            typer.echo("Looking up channel...")
            element = WebDriverWait(driver=driver, timeout=5).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, sel)))
        if driver.find_element(by=By.CSS_SELECTOR, value=sel):
            is_live = True
    except TimeoutException:
        is_live = False

    driver.close()

    return func_return(is_live, should_print)


@app.command()
def get_categories(
    should_print: bool = typer.Option(False, '--print', '-p'),
    pretty: bool = typer.Option(False, '--pretty', '-t')
):
    """
    Gets top 30 Twitch categories.
    SHOULDPRINT determines if result is returned to stdout
    or to function caller.
    PRETTY determines if output is pretty printed or not.
    """
    URL = "https://www.twitch.tv/directory"

    driver = get_driver()
    driver.get(URL)

    with click_spinner.spinner():
        typer.echo("Looking up categories...")
        element = WebDriverWait(driver=driver, timeout=5).until(
            EC.presence_of_element_located(
                (By.CSS_SELECTOR, 'div[data-target="directory-first-item"]'))
        )

        # init results string
        results = ""

        # get categories first entry
        firstEl = driver.find_element(
            by=By.CSS_SELECTOR, value='div[data-target="directory-first-item"]')
        results += firstEl.find_element(by=By.CSS_SELECTOR,
                                        value='h2[title]').text + SEP
        results += firstEl.find_element(by=By.CSS_SELECTOR,
                                        value='p a').text.replace('viewers', '') + SEP
        results += firstEl.find_element(
            by=By.CSS_SELECTOR,
            value='a[data-a-target="tw-box-art-card-link"]'
        ).get_attribute('href') + '\n'

        # get categories other entries
        for item in driver.find_elements(by=By.CSS_SELECTOR, value='div[data-target=""]'):
            results += item.find_element(by=By.CSS_SELECTOR,
                                         value='h2[title]').text + SEP
            results += item.find_element(by=By.CSS_SELECTOR,
                                         value='p a').text.replace('viewers', '') + SEP
            results += item.find_element(
                by=By.CSS_SELECTOR,
                value='a[data-a-target="tw-box-art-card-link"]'
            ).get_attribute('href') + '\n'

    driver.close()

    return func_return(results, should_print, pretty)


@app.command()
def get_channels(
    target: str,
    should_print: bool = typer.Option(False, '--print', '-p'),
    pretty: bool = typer.Option(False, '--pretty', '-t')
):
    """
    Gets live channels in TARGET, it can be a channel link, or just channel name.
    SHOULDPRINT determines if result is returned to stdout or function caller.
    PRETTY determines if output is pretty printed or not.
    """
    if target.find("https://") > -1:
        URL = f'{target}?sort=VIEWER_COUNT'
    else:
        URL = f'https://www.twitch.tv/directory/game/{target}?sort=VIEWER_COUNT'

    driver = get_driver()
    driver.get(URL)

    with click_spinner.spinner():
        typer.echo("Lookin up channels...")
        element = WebDriverWait(driver=driver, timeout=5).until(
            EC.presence_of_element_located(
                (By.CSS_SELECTOR, 'div[data-target="directory-container"]'))
        )

        el = driver.find_element(
            by=By.CSS_SELECTOR, value='div[data-target="directory-container"]')

        # init results
        results = ""

        # get channels first entry
        firstEl = el.find_element(
            by=By.CSS_SELECTOR, value='div[data-target="directory-first-item"]')
        results += firstEl.find_element(
            by=By.CSS_SELECTOR,
            value='a[data-a-target="preview-card-channel-link"]'
        ).get_attribute('href').split('/')[-2] + SEP
        results += firstEl.find_element(by=By.CSS_SELECTOR,
                                        value='h3[title]').text + '\n'

        # get channels other entries
        for item in el.find_elements(by=By.CSS_SELECTOR, value='div[data-target=""]'):
            results += item.find_element(
                by=By.CSS_SELECTOR,
                value='a[data-a-target="preview-card-channel-link"]'
            ).get_attribute('href').split('/')[-2] + SEP
            results += item.find_element(by=By.CSS_SELECTOR,
                                         value='h3[title]').text + '\n'

    driver.close()

    return func_return(results, should_print, pretty)


@app.command()
def choose(pretty: bool = typer.Option(False, '--pretty', '-t')):
    """
    Gets top 30 Twitch categories, requests user to choose one
    and then gets current live channels.
    PRETTY determines if output is pretty printed or not.
    """
    categories = get_categories(should_print=False)

    splitCategories = []
    for cat in categories.split('\n'):
        splited = cat.split(SEP)
        if len(splited) < 3:
            continue
        catObj = {
            "name": splited[0].strip(),
            "views": splited[1].strip(),
            "link": splited[2].strip()
        }
        splitCategories.append(catObj)

    if pretty:
        print(tabulate(splitCategories, showindex="always",
              headers="keys", tablefmt="plain"))
    else:
        for i, cat in enumerate(splitCategories):
            print(i, cat["name"], cat["views"])

    try:
        choosen = int(input("\nChoose a category by index number:"))
        if choosen > len(splitCategories) or choosen < 0:
            raise Exception()
    except:
        typer.secho("ERROR: Invalid index",
                    fg=typer.colors.WHITE, bg=typer.colors.RED)
        raise typer.Exit(code=1)

    get_channels(splitCategories[choosen]["link"],
                 should_print=True, pretty=pretty)


if __name__ == "__main__":
    app()