Data-Analysis-Jupyter/medium/retrieval.py

from multiprocessing import Pool
import requests
import re
from bs4 import BeautifulSoup
from itertools import chain
from collections import Counter
from timeit import default_timer as timer
import pandas as pd
from datetime import datetime


def get_table_rows(fname='stats.html'):
    """
    Extract the table rows from the statistics

    :param fname: string name of the file stored in `data` directory

    :return table_rows: list of BeautifulSoup objects to be passed to `process_in_parallel`
    """

    soup = BeautifulSoup(open(f'data/{fname}', 'r'), features='lxml')
    table_rows = soup.find_all(
        attrs={'class': "sortableTable-row js-statsTableRow"})
    print(f'Found {len(table_rows)} entries in table.')
    return table_rows


def convert_timestamp(ts: int, tz: str):
    """Convert a unix timestamp to a date timestamp"""
    return pd.to_datetime(ts, origin='unix', unit='ms').tz_localize('UTC').tz_convert(tz).tz_localize(None)


def process_entry(entry, parallel=True, tz='America/Chicago'):
    """
    Extract data from one entry in table

    :param entry: BeautifulSoup tag
    :param parallel: Boolean for whether function is being run in parallel
    :param tz: string representing timezone for started and published time

    :return entry_dict: dictionary with data about entry

    """
    # Convert to soup when running in parallel
    if parallel:
        entry = BeautifulSoup(entry, features='lxml').body.tr

    entry_dict = {}
    # Extract information
    for value, key in zip(entry.find_all(attrs={'class': 'sortableTable-value'}),
                          ['published_date', 'views', 'reads', 'ratio', 'fans']):
        entry_dict[key] = float(
            value.text) if key == 'ratio' else int(value.text)

    entry_dict['read_time'] = int(entry.find_all(attrs={'class': 'readingTime'})[
                                  0].get('title').split(' ')[0])

    # Unlisted vs published
    entry_dict['type'] = 'unlisted' if len(
        entry.find_all(text=' Unlisted')) > 0 else 'published'

    # Publication
    publication = entry.find_all(attrs={'class': 'sortableTable-text'})
    if 'In' in publication[0].text:
        entry_dict['publication'] = publication[0].text.split('In ')[
            1].split('View')[0]
    else:
        entry_dict['publication'] = 'None'

    # Convert datetimes
    entry_dict['published_date'] = convert_timestamp(
        entry_dict['published_date'], tz=tz)
    entry_dict['started_date'] = convert_timestamp(
        entry.get('data-timestamp'), tz=tz)

    # Get the link
    link = entry.find_all(text='View story',
                               attrs={'class': 'sortableTable-link'})[0].get('href')
    entry_dict['link'] = link
    # Retrieve the article and create a soup
    entry = requests.get(link).content
    entry_soup = BeautifulSoup(entry, features='lxml')

    # Get the title
    try:
        title = entry_soup.h1.text
    except:
        title = 'response'

    title_word_count = len(re.findall(r"[\w']+|[.,!?;]", title))

    # Main text entries
    entry_text = [p.text for p in entry_soup.find_all(
        ['h1', 'h2', 'h3', 'p', 'blockquote'])]

    # Make sure to catch everything
    entry_text.extend(s.text for s in entry_soup.find_all(
        attrs={'class': 'graf graf--li graf-after--li'}))
    entry_text.extend(s.text for s in entry_soup.find_all(
        attrs={'class': 'graf graf--li graf-after--p'}))
    entry_text.extend(s.text for s in entry_soup.find_all(
        attrs={'class': 'graf graf--li graf-after--blockquote'}))
    entry_text.extend(s.text for s in entry_soup.find_all(
        attrs={'class': 'graf graf--li graf-after--pullquote'}))

    entry_text = ' '.join(entry_text)

    # Word count
    word_count = len(re.findall(r"[\w']+|[.,!?;]", entry_text))

    # Number of claps
    clap_pattern = re.compile(
        '^[0-9]{1,} claps|^[0-9]{1,}.[0-9]{1,}K claps|^[0-9]{1,}K claps')
    claps = entry_soup.find_all(text=clap_pattern)

    if len(claps) > 0:
        if 'K' in claps[0]:
            clap_number = int(1e3 * float(claps[0].split('K')[0]))
        else:
            clap_number = int(claps[0].split(' ')[0])
    else:
        clap_number = 0

    # Post tags
    tags = entry_soup.find_all(
        attrs={'class': 'tags tags--postTags tags--borderless'})
    tags = [li.text for li in tags[0].find_all('li')]

    # Responses to entry
    responses = entry_soup.find_all(attrs={'class': 'button button--chromeless u-baseColor--buttonNormal u-marginRight12',
                                           'data-action': 'scroll-to-responses'})
    num_responses = int(responses[0].text) if len(responses) > 0 else 0

    # Store in dictionary
    entry_dict['title'] = title
    entry_dict['title_word_count'] = title_word_count
    entry_dict['text'] = entry_text
    entry_dict['word_count'] = word_count
    entry_dict['claps'] = clap_number
    entry_dict['tags'] = tags
    entry_dict['num_responses'] = num_responses

    # Time since publication
    entry_dict['days_since_publication'] = (
        datetime.now() - entry_dict['published_date']).total_seconds() / (3600 * 24)

    return entry_dict


def process_in_parallel(table_rows, processes=20):
    """
    Process all the stats in a table in parallel

    :note: make sure to set the correct time zone in `process_entry`
    :note: running on Mac may first require setting
    export OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES
    from the command line to enable parallel processing

    :param table_rows: BeautifulSoup table rows

    :param processes: integer number of processes (threads) to use in parallel

    :return df: dataframe of information about each post

    """
    # Convert to strings for multiprocessing
    table_rows_str = [str(r) for r in table_rows]

    # Process each article in paralllel
    pool = Pool(processes=processes)
    results = []
    start = timer()
    for i, r in enumerate(pool.imap_unordered(process_entry, table_rows_str)):
        # Report progress
        print(f'{100 * i / len(table_rows_str):.2f}% complete.', end='\r')
        results.append(r)
    pool.close()
    pool.join()
    end = timer()
    print(
        f'Processed {len(table_rows_str)} articles in {end-start:.2f} seconds.')

    # Convert to dataframe
    df = pd.DataFrame(results)
    # Rename ratio
    df.rename(columns={'ratio': 'read_ratio'}, inplace=True)
    # Add extra columns with more data
    df['claps_per_word'] = df['claps'] / df['word_count']
    df['editing_days'] = ((df['published_date'] - df['started_date']
                           ).dt.total_seconds() / (60 * 60 * 24)).astype(int)

    # Rounding
    df['published_date'] = df['published_date'].dt.round('min')
    df['started_date'] = df['started_date'].dt.round('min')
    df['read_ratio'] = df['read_ratio'].round(2)

    # 5 most common tags (might want to include more tags)
    n = 5
    all_tags = list(chain(*df['tags'].tolist()))
    tag_counts = Counter(all_tags)
    tags = tag_counts.most_common(n)

    # Adding columns with indication of tag
    for tag, count in tags:
        flag = [1 if tag in tags else 0 for tags in df['tags']]
        df.loc[:, f'<tag>{tag}'] = flag

    df.sort_values('published_date', inplace=True)
    return df


def get_data(fname='stats.html', processes=20):
    """
    Retrieve medium article statistics

    :note: running on Mac may first require setting
    export OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES
    from the command line to enable parallel processing

    :param fname: file name (should be 'stats.html')
    :param processes: integer number of processes

    :return df: dataframe of article data
    """
    t = get_table_rows(fname=fname)
    return process_in_parallel(t, processes=processes)