Stock data REST API extractor with GCP storage options

Question

I'm currently working on an AI-driven trading system, the code below aims to extract ticker data from polygon REST API, this is a paid service so, in order to test the code you will need to subscribe / obtain a free API key with limited data history. You'll find base_extractor.py, polygon_extractor.py and extract.py which I will explain briefly above each.

My main concerns:

• I'm concerned with the intraday data (1min or less) for technical reasons, those who are experienced with trading will understand its significance. Anyway the API limits the number of records(minute price data point in this case) to 5000 minutes max per GET request, therefore you'll come across a parameter called days_per_request which main purpose is to control the rate of records returned per request. Of course this negatively impacts the time requirements so any suggestions to improve this bottleneck, will greatly impact the efficiency of the extractor.
• Modularization issues that I overcome with sys.path.append('..') which I need to get rid of without PyCharm complaining about unresolved references that resolve somehow by runtime. You will understand further if you read through the code.
• General suggestions and feedback about the whole code as well as performance / speed improvements / general structure are more than welcome.
• Is using concurrent.futures for sending concurrent http requests the best option? or do you have other suggestions that are faster?

base_extractor.py: the base class that contains methods that are common to this extraction process regardless of the API and can be used with polygon and for other REST APIs that provide the same service(most of them have the same design). It contains useful features including memoryless writing of data to .parquet format and storing to GCP cloud storage(optional).

from oauth2client.service_account import ServiceAccountCredentials
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from logging import handlers
import pyarrow.parquet as pq
from gcloud import storage
import pyarrow as pa
import pandas as pd
import requests
import logging
import shutil
import json
import os


class BaseExtractor:
    """
    A tool for downloading stock data from these websites:
        - https://www.tiingo.com
        - https://www.polygon.io
    """

    def __init__(
        self,
        api_key,
        base_url,
        compression='gzip',
        log_file=None,
        workers=4,
        single_file=False,
        gcp_bucket=None,
        gcp_key=None,
        request_headers=None,
    ):
        """
        Initialize extractor
        Args:
            api_key: Key provided by the target website.
            base_url: API base url.
            compression:
                parquet compression types:
                    - 'brotli'
                    - 'snappy'
                    - 'gzip'
            log_file: Path to log file.
            workers: Concurrent connections.
            single_file: Single file per extraction.
            gcp_bucket: Google bucket name.
            gcp_key: Google bucket authentication json key file.
            request_headers: HTTP headers that will be used with requests.
        """
        self.api_key = api_key
        self.base_url = base_url
        self.compression = compression
        self.log_file_name = log_file
        self.logger = self.get_logger()
        self.workers = workers
        self.single_file = single_file
        self.gcp_bucket = gcp_bucket
        self.gcp_key = gcp_key
        self.request_headers = request_headers

    def write_results(self, response, fp, json_key=None):
        """
        Write extractions to a supported format [.parquet]
        Args:
            response: API response.
            fp: Path to output file.
            json_key: Key in response.json()

        Returns:
            None
        """
        if results := (response.json().get(json_key) if json_key else response.json()):
            frame = pd.DataFrame(results)
            frame[frame.T.dtypes == int] = frame[frame.T.dtypes == int].astype(float)
            if fp.endswith('.parquet'):
                table = pa.Table.from_pandas(frame)
                pq.write_to_dataset(table, root_path=fp, compression=self.compression)

    def get_logger(self):
        """
        Create logger.

        Returns:
            logger object.
        """
        formatter = logging.Formatter(
            '%(asctime)s %(name)s: ' '%(levelname)-2s %(message)s'
        )
        logger = logging.getLogger('API Extractor')
        logger.setLevel(logging.DEBUG)
        if self.log_file_name:
            file_handler = handlers.RotatingFileHandler(
                self.log_file_name, backupCount=10
            )
            file_handler.setFormatter(formatter)
            logger.addHandler(file_handler)
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)
        return logger

    def extract_data(self, method, urls, *args, **kwargs):
        """
        Extract urls from a supported API.
        Args:
            method: One of BaseExtractor extraction methods.
            urls: A list of full urls that will be extracted by the given method.
            *args: method args.
            **kwargs: method kwargs.

        Returns:
            None
        """
        with ThreadPoolExecutor(max_workers=self.workers) as executor:
            future_requests = {
                executor.submit(method, url, *args, **kwargs): url for url in urls
            }
            for future_response in as_completed(future_requests):
                try:
                    future_response.result()
                except Exception as e:
                    self.logger.exception(
                        f'Failed to get {future_requests[future_response]}\n{e}'
                    )

    @staticmethod
    def get_intervals(
        start_date, end_date=None, days_per_request=5, date_fmt='%Y-%m-%d'
    ):
        """
        Get all date intervals that need to be extracted.
        Args:
            start_date: Timestamp / datetime.
            end_date: Timestamp / datetime.
            days_per_request: Maximum days per HTTP request.
            date_fmt: Output interval date format.

        Returns:
            start_intervals, end_intervals
        """
        start_intervals = pd.date_range(
            start_date,
            end_date or datetime.now(),
            freq=f'{days_per_request + 1}d',
        )
        end_intervals = start_intervals + pd.offsets.Day(days_per_request)
        return [
            interval.to_series().dt.strftime(date_fmt)
            for interval in (start_intervals, end_intervals)
        ]

    def store_gcp_bucket(self, fp):
        """
        Store data to google bucket.
        Args:
            fp: Filepath to be stored(folder or file).

        Returns:
            None
        """
        gcp_credentials = None
        if self.gcp_key:
            with open(self.gcp_key) as key:
                gcp_credentials = json.load(key)
                gcp_credentials = ServiceAccountCredentials.from_json_keyfile_dict(
                    gcp_credentials
                )
        client = storage.Client(credentials=gcp_credentials)
        bucket = client.get_bucket(self.gcp_bucket)
        self.upload_to_gcp(fp, bucket)

    def upload_to_gcp(self, fp, bucket):
        """
        Upload a given filepath to GCP bucket.
        Args:
            fp: Filepath to be uploaded(folder or file).
            bucket: gcloud.storage.bucket.Bucket

        Returns:
            None
        """
        if os.path.isfile(fp):
            blob = bucket.blob(fp)
            blob.upload_from_filename(fp)
            self.delete_file(fp)
            self.logger.info(f'Transfer of gs://{fp} complete')
        if os.path.isdir(fp):
            fps = [os.path.join(fp, f) for f in os.listdir(fp)]
            for fp in fps:
                self.upload_to_gcp(fp, bucket)

    def finalize_extraction(self, fp, sort_column=None):
        """
        Process file after extraction.
        Args:
            fp: Path to output file.
            sort_column: Column to sort data by.

        Returns:
            None
        """
        if not os.path.exists(fp):
            self.logger.info(f'Expected to find {fp}')
            return
        if self.single_file:
            temp = pd.read_parquet(fp)
            self.delete_file(fp)
            if sort_column and sort_column in temp.columns:
                temp = temp.set_index(sort_column).sort_index()
            temp.to_parquet(fp)
        if self.gcp_bucket:
            self.store_gcp_bucket(fp)

    @staticmethod
    def join_query(query_args, **kwargs):
        """
        Join query args.
        Args:
            query_args: A dictionary that contains args and their values.
            **kwargs: Additional args and their values.

        Returns:
            joined query.
        """
        query_args.update(kwargs)
        return '&'.join(f'{arg}={val}' for arg, val in query_args.items())

    @staticmethod
    def delete_file(fp):
        """
        Delete a file from disk.
        Args:
            fp: Path to file to be deleted.

        Returns:
            None
        """
        if os.path.isdir(fp):
            shutil.rmtree(fp)
        if os.path.isfile(fp):
            os.remove(fp)

    def get_url(self, full_url):
        """
        Send a GET request.
        Args:
            full_url: Full url with target args.

        Returns:
            response.
        """
        response = requests.get(full_url, headers=self.request_headers)
        self.logger.info(f'Got response {response} for {full_url}')
        return response

polygon_extractor.py is BaseExtractor subclass and has methods specific to polygon API. You will come across sys.path.append() I mentioned earlier that I need to replace without introducing issues to the code. extractors is the name of the enclosing repo subfolder that contains extraction modules.

import sys

sys.path.append('..')

from extractors.base_extractor import BaseExtractor
from collections import defaultdict
from pathlib import Path


class PolygonExtractor(BaseExtractor):
    """
    A tool for downloading data from polygon.io API
    """

    def __init__(
        self,
        api_key,
        base_url='https://api.polygon.io',
        compression='gzip',
        log_file=None,
        workers=4,
        single_file=False,
        gcp_bucket=None,
        gcp_key=None,
    ):
        """
        Initialize extractor
        Args:
            api_key: Key provided by polygon.io API.
            base_url: https://api.polygon.io
            compression:
                parquet compression types:
                    - 'brotli'
                    - 'snappy'
                    - 'gzip'
            log_file: Path to log file.
            workers: Concurrent connections.
            single_file: Single file per extraction.
            gcp_bucket: Google bucket name.
            gcp_key: Google bucket authentication json key file.
        """
        self.ticker_extraction_counts = defaultdict(lambda: 0)
        super(PolygonExtractor, self).__init__(
            api_key,
            base_url,
            compression,
            log_file,
            workers,
            single_file,
            gcp_bucket,
            gcp_key,
        )

    def extract_agg_page(self, full_url, ticker, interval, fp):
        """
        Extract a single page ticker data from urls with the following prefix:
        https://api.polygon.io/v2/aggs/ticker/

        Args:
            full_url: Full url with the valid prefix and args.
            ticker: One of the tickers supported ex: 'AAPL'
            interval: One of the following:
                - 'minute'
                - 'hour'
                - 'day'
                - 'week'
                - 'month'
                - 'quarter'
                - 'year'
            fp: Path to output file.

        Returns:
            None
        """
        response = self.get_url(full_url)
        start_date, end_date = full_url.split('/')[10:12]
        self.logger.info(
            f'Extracted {ticker} aggregate {interval} data '
            f'[{start_date}] --> [{end_date[:10]}] | url: {full_url}'
        )
        self.write_results(response, fp, 'results')

    def extract_ticker_page(self, full_url, market, fp, total_pages=1):
        """
        Extract a single page ticker data from urls with the following prefix.
        https://api.polygon.io/v2/reference/tickers

        Args:
            full_url: Full url with the valid prefix.
            market: One of the supported markets.
            fp: Path to output file.
            total_pages: Total number of pages that are being extracted.

        Returns:
            None
        """
        response = self.get_url(full_url)
        self.ticker_extraction_counts[market] += 1
        completed = self.ticker_extraction_counts[market]
        self.logger.info(
            f'Extracted {market} ticker page: {completed}/{total_pages} url: {full_url}'
        )
        self.write_results(response, fp, 'tickers')

    def extract_available_tickers(
        self,
        fp,
        sort_by='ticker',
        market='STOCKS',
        per_page=2000,
        sort_column=None,
        **kwargs,
    ):
        """
        Extract all available tickers for a given market
        Args:
            fp: Path to output file
            sort_by: 'ticker' or 'type'
            market: One of the following options:
                - 'STOCKS'
                - 'INDICES'
                - 'CRYPTO'
                - 'FX'
            per_page: Results returned per result page
            sort_column: Column name to use for sorting the data.
            **kwargs: Additional query args

        Returns:
            None
        """
        self.logger.info(f'Started extraction of {market} available tickers')
        query_args = {
            'sort': sort_by,
            'market': market,
            'perpage': per_page,
            'page': '1',
        }
        query_args = self.join_query(query_args, **kwargs)
        query_contents = [
            self.base_url,
            'v2',
            'reference',
            f'tickers?{query_args}&apiKey={self.api_key}',
        ]
        full_link = '/'.join(query_contents)
        count = int(self.get_url(full_link).json()['count'])
        page_count = (count // per_page) + 1
        target_urls = [
            full_link.replace('page=1', f'page={i}') for i in range(1, page_count + 1)
        ]
        self.extract_data(self.extract_ticker_page, target_urls, market, fp, page_count)
        self.finalize_extraction(fp, sort_column)
        self.logger.info(f'Finished extraction of {market} available tickers')

    def extract_ticker(
        self,
        fp,
        ticker,
        start_date,
        end_date=None,
        days_per_request=5,
        interval='day',
        multiplier='1',
        date_fmt='%Y-%m-%d',
        sort_column=None,
        **kwargs,
    ):
        """
        Extract data of a supported ticker for a specified period of time
        Args:
            fp: Path to output file
            ticker: A supported ticker ex: 'AAPL'
            start_date: A date in the following format yy-mm-dd to start from
            end_date: A date in the following format yy-mm-dd to stop at
            days_per_request: Days to extract per get request
            interval: interval between data points, options are:
                - 'minute'
                - 'hour'
                - 'day'
                - 'week'
                - 'month'
                - 'quarter'
                - 'year'
            multiplier: Size of the timespan multiplier
            date_fmt: Date interval format, default yy-mm-dd
            sort_column: Column name to use for sorting the data.
            **kwargs: Additional query args.

        Returns:
            None
        """
        self.logger.info(f'Started extraction of {ticker}')
        start_intervals, end_intervals = self.get_intervals(
            start_date, end_date, days_per_request, date_fmt
        )
        query_args = self.join_query({}, **kwargs)
        query_contents = [
            self.base_url,
            'v2',
            'aggs',
            'ticker',
            ticker,
            'range',
            multiplier,
            interval,
            'start_date',
            f'end_date?{query_args}&apiKey={self.api_key}',
        ]
        full_url = '/'.join(query_contents)
        target_urls = [
            full_url.replace('start_date', d1).replace('end_date', d2)
            for d1, d2 in zip(start_intervals, end_intervals)
        ]
        self.extract_data(self.extract_agg_page, target_urls, ticker, interval, fp)
        self.finalize_extraction(fp, sort_column)
        self.logger.info(f'Finished extraction of {ticker}')

    def extract_tickers(self, ticker_file, destination='.', *args, **kwargs):
        """
        Extract ticker data from a file containing a list of tickers.
        Args:
            ticker_file: Filepath that contains target tickers.
            destination: Path to destination folder.
            *args: self.extract_ticker() args.
            **kwargs: self.extract_ticker() kwargs.

        Returns:
            None
        """
        tickers = [item for item in open(ticker_file)]
        total = len(tickers)
        for i, ticker in enumerate(tickers):
            fp = Path(destination) / Path(f'{(ticker := ticker.strip())}.parquet')
            self.extract_ticker(str(fp), ticker, *args, **kwargs)
            self.logger.info(
                f'Extracted {i + 1}/{total} tickers | '
                f'completed: {100 * ((i + 1) / total)}%'
            )

extract.py is the cli parsing module that defines general as well as API specific args. And it allows control over the whole extraction operation from the command line.

#!/usr/local/bin/python3.8
import argparse
import sys

sys.path.append('..')

from extractors.polygon_extractor import PolygonExtractor
from extractors.tiingo_extractor import TiingoExtractor
import os
import sys


def process_polygon(cli_args, extractor):
    """
    Perform extraction through polygon.io API
    Args:
        cli_args: Command line args.
        extractor: BaseExtractor subclass.

    Returns:
        None
    """
    if cli_args.available:
        extractor.extract_available_tickers(
            cli_args.output,
            market=cli_args.market,
            per_page=cli_args.per_page,
            sort_column=cli_args.sort_column,
        )
    if cli_args.ticker:
        assert cli_args.ticker, f'ticker not specified'
        assert cli_args.start_date, f'start date not specified'
        assert cli_args.output, f'Output file not specified'
        extractor.extract_ticker(
            cli_args.output,
            cli_args.ticker,
            cli_args.start_date,
            cli_args.end_date,
            cli_args.days_per_request,
            cli_args.interval,
            sort_column=cli_args.sort_column,
        )
    if cli_args.tickers:
        os.makedirs(cli_args.output, exist_ok=True)
        extractor.extract_tickers(
            cli_args.tickers,
            cli_args.output,
            cli_args.start_date,
            cli_args.end_date,
            cli_args.days_per_request,
            cli_args.interval,
            sort_column=cli_args.sort_column,
        )


def process_from_cli(parser, argv):
    """
    Parse cli args and initialize extractor.
    Args:
        parser: argparse.ArgumentParser()
        argv: sys.argv

    Returns:
        None
    """
    extractors = {'tiingo': TiingoExtractor, 'polygon': PolygonExtractor}
    cli_args = parser.parse_args(argv)
    assert (target := cli_args.target) in extractors, 'unsupported api'
    extractor = extractors[target](
        api_key=cli_args.key,
        compression=cli_args.compression,
        log_file=cli_args.log,
        workers=cli_args.workers,
        single_file=cli_args.single_file,
        gcp_bucket=cli_args.gcp_bucket,
        gcp_key=cli_args.gcp_key,
    )
    if target == 'polygon':
        process_polygon(cli_args, extractor)


def default_args():
    """
    Define default cli args that are common between supported APIs.

    Returns:
        parser, extraction_group
    """
    parser = argparse.ArgumentParser()
    extraction_group = parser.add_mutually_exclusive_group()
    extraction_group.add_argument('--ticker', help="a single ticker ex: 'AAPL'")
    extraction_group.add_argument('--tickers', help='a file that contains tickers')
    parser.add_argument('-k', '--key', help='polygon.io api key', required=True)
    parser.add_argument(
        '-t', '--target', help="One of the supported apis ex: 'tiingo'", required=True
    )
    parser.add_argument(
        '-o', '--output', help='path to a file or folder', required=True
    )
    parser.add_argument(
        '-c', '--compression', help='compression type', default='brotli'
    )
    parser.add_argument('-l', '--log', help='log file path')
    parser.add_argument(
        '-w', '--workers', help='concurrent requests', default=4, type=int
    )
    parser.add_argument(
        '--single_file',
        action='store_true',
        help='combine .parquet file chunks in a single file',
    )
    parser.add_argument(
        '--start_date', help="start date of extraction for timed data ex: '2020-01-30'"
    )
    parser.add_argument(
        '--end_date', help='end date of extraction for timed data', default=None
    )
    parser.add_argument(
        '--gcp_key', help='Google cloud json authentication file', default=None
    )
    parser.add_argument('--gcp_bucket', help='Google cloud bucket name', default=None)
    parser.add_argument(
        '--days_per_request',
        help='day interval per get request',
        default=5,
        type=int,
    )
    parser.add_argument(
        '--interval', help='interval between data points', default='day'
    )
    parser.add_argument(
        '--sort_column', help='column name to sort data by', default=None
    )
    return parser, extraction_group


def get_polygon_args(parser, extraction_group):
    """
    Define args that are specific to polygon.io API.
    Args:
        parser: argparse.ArgumentParser()
        extraction_group: Extraction mutually exclusive group.

    Returns:
        parser
    """
    extraction_group.add_argument(
        '--available', action='store_true', help='extract available tickers'
    )
    parser.add_argument('--market', help='market to extract', default='STOCKS')
    parser.add_argument(
        '--per_page', help='records per response page', default=2000, type=int
    )
    return parser


def tiingo_args():
    pass


def main(argv):
    parser, extraction_group = default_args()
    updated_parser = get_polygon_args(parser, extraction_group)
    process_from_cli(updated_parser, argv)


if __name__ == '__main__':
    main(sys.argv[1:])

Answer 1

The biggest thing that stands out to me is the repetition in your code. The same large groups of variables are written out and passed around in the same order repeatedly and the same function is called over and over for different arguments. Those are signs that what you're doing should probably be simplified.

In particular the model of your config and tickers can be more clearly separated from the actions you use them for.

For example, BaseExtractor and PolygonExtractor repeat the same 9 variables 5 separate times between being used as paramaters and values to set. That could be cut down to once with dataclasses and multiple inheritance:

from dataclasses import dataclass
from collections import defaultdict


@dataclass 
class BaseExtractorConfig:
    api_key:                str
    base_url:               str
    compression:            str ='gzip'
    log_file:               str = None
    workersL:               int = 4
    single_file:            bool = False
    gcp_bucket:             str = None
    gcp_key:                str = None
    request_headers:        str = None
    logger:                 str = None

    def __post_init__(self):
        self.logger = self.get_logger()


class BaseExtractor(BaseExtractorConfig):   
    def get_logger(self):
        return 'logger set'


@dataclass
class PolygonExtractorConfig(BaseExtractorConfig):
    base_url:                   str = 'https://api.polygon.io'
    ticker_extraction_counts:   dict = None
    
    def __post_init__(self):
        super().__post_init__()
        self.ticker_extraction_counts = defaultdict(lambda: 0)


class PolygonExtractor(PolygonExtractorConfig, BaseExtractor):
    def f(self):
        print(self)

pe = PolygonExtractor('api_key_here', gcp_key="added a kwargs")
pe.f()

which prints

PolygonExtractor(api_key='api_key_here', base_url='https://api.polygon.io', compression='gzip', log_file=None, workersL=4, single_file=False, gcp_bucket=None, gcp_key='added a kwargs', request_headers=None, logger='logger set', ticker_extraction_counts=defaultdict(<function PolygonExtractorConfig.__post_init__.<locals>.<lambda> at 0x7f43344e73a0>, {}))

You could take a similar approach to the ticker values which would make it much easier to follow what are objects being used in your code and what are actions being performed.

I would also split the parser arguments into a separate json file or the like, read them in as a list, and then add them all with a single loop. The external file would more clearly show the commands and their structures while the code in python would be cleaner.