{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-01-05T05:06:16.260Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import sys\n",
    "import re\n",
    "import datetime\n",
    "import dateutil\n",
    "import numpy as np\n",
    "import concurrent\n",
    "import pyarrow as pa\n",
    "import pathlib\n",
    "from timeit import default_timer as timer\n",
    "\n",
    "from pyqstrat import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2019-01-05T05:01:48.768Z"
    },
    "run_control": {
     "frozen": false,
     "read_only": false
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "VERBOSE = False\n",
    "\n",
    "class PathFileNameProvider:\n",
    "    \"\"\"A helper class that, given a pattern such as such as \"/tmp/abc*.gz\" and an optional include and exclude pattern, \n",
    "    returns names of all files that match\n",
    "    \"\"\"\n",
    "    def __init__(self, path, include_pattern = None, exclude_pattern = None):\n",
    "        '''\n",
    "        Args:\n",
    "            path (str): A pattern such as \"/tmp/abc*.gz\"\n",
    "            include_pattern (str): Given a pattern such as \"xzy\", will return only filenames that contain xyz\n",
    "            exclude_pattern (str): Given a pattern such as \"_tmp\", will exclude all filenames containing _tmp\n",
    "        '''\n",
    "        self.path = path\n",
    "        self.include_pattern = include_pattern\n",
    "        self.exclude_pattern = exclude_pattern\n",
    "        \n",
    "    def __call__(self):\n",
    "        \"\"\"\n",
    "        Returns:\n",
    "            list of str: all matching filenames\n",
    "        \"\"\"\n",
    "        files = glob.glob(self.path)\n",
    "        if not len(files):\n",
    "            raise Exception(f'no matching files found with pattern: {self.path}')\n",
    "        if self.include_pattern:\n",
    "            files = [file for file in files if self.include_pattern in file]\n",
    "        if self.exclude_pattern:\n",
    "            files = [file for file in files if self.exclude_pattern not in file]\n",
    "        if not len(files):\n",
    "            raise Exception(f'no matching files for: {self.path} including: {self.include_pattern} excluding: {self.exclude_pattern}')\n",
    "        return files\n",
    "    \n",
    "class SingleDirectoryFileNameMapper:\n",
    "    \"\"\"A helper class that provides a mapping from input filenames to their corresponding output filenames in an output directory.\"\"\"\n",
    "    def __init__(self, output_dir):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            output_dir (str): The directory where we want to write output files\n",
    "        \"\"\"\n",
    "        if not os.path.isdir(output_dir): raise Exception(f'{output_dir} does not exist')\n",
    "        self.output_dir = output_dir\n",
    "\n",
    "    def __call__(self, input_filepath):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            input_filepath (str): The input file that we are creating an output file for, e.g. \"/home/xzy.gz\"\n",
    "        Returns:\n",
    "            str: Output file path for that input.  We take the filename from the input filepath, strip out any extension \n",
    "                and prepend the output directory name\n",
    "        \"\"\"\n",
    "        \n",
    "        if self.output_dir is None:\n",
    "            dirname = os.path.dirname(input_filepath)\n",
    "            dirname = os.path.join(dirname, 'output')\n",
    "        else:\n",
    "            dirname = self.output_dir\n",
    "            \n",
    "        if not os.path.isdir(dirname): raise Exception(f'{dirname} does not exist')\n",
    "     \n",
    "        input_filename = os.path.basename(input_filepath)\n",
    "        exts = r'\\.txt$|\\.gz$|\\.bzip2$|\\.bz$|\\.tar$|\\.zip$|\\.csv$'\n",
    "        while (re.search(exts, input_filename)):\n",
    "            input_filename = '.'.join(input_filename.split('.')[:-1])\n",
    "            if VERBOSE: print(f'got input file: {input_filename}')\n",
    "        output_prefix = os.path.join(dirname, input_filename)\n",
    "        return output_prefix\n",
    "\n",
    "class TextHeaderParser:\n",
    "    \"\"\"\n",
    "    Parses column headers from a text file containing market data\n",
    "    \"\"\"\n",
    "    def __init__(self, record_generator, skip_rows = 0, separator = ',', make_lowercase = True):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "        \n",
    "            record_generator: A function that takes a filename and its compression type and returns an object\n",
    "                that we can use to iterate through lines in that file\n",
    "            skip_rows (int, optional): Number of rows to skip before starting to read the file.  Default is 0\n",
    "            separator (str, optional): Separator for headers.  Defaults to ,\n",
    "            make_lowercase (bool, optional): Whether to convert headers to lowercase before returning them\n",
    "        \"\"\"\n",
    "        self.record_generator = record_generator\n",
    "        self.skip_rows = 0\n",
    "        self.separator = separator\n",
    "        self.make_lowercase = make_lowercase\n",
    "        \n",
    "    def __call__(self, input_filename, compression):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "        \n",
    "        input_filename (str): The file to read\n",
    "        compression (str): Compression type, e.g. \"gzip\", or None if the file is not compressed\n",
    "        \n",
    "        Returns:\n",
    "            list of str: column headers\n",
    "        \"\"\"\n",
    "        \n",
    "        decode_needed = (compression is not None and compression != \"\")\n",
    "        \n",
    "        with self.record_generator(input_filename, compression) as f:\n",
    "            headers = None\n",
    "            for line_num, line in enumerate(f):\n",
    "                if decode_needed: line = line.decode()\n",
    "                if line_num < self.skip_rows: continue\n",
    "                headers = line.split(self.separator)\n",
    "                headers = [re.sub('[^A-Za-z0-9 ]+', '', header) for header in headers]\n",
    "                if len(headers) == 1:\n",
    "                    raise Exception(f'Could not parse headers: {line} with separator: {self.separator}')\n",
    "                break\n",
    "\n",
    "            if headers is None: raise Exception('no headers found')\n",
    "            if self.make_lowercase: headers = [header.lower() for header in headers]\n",
    "            if VERBOSE: print(f'Found headers: {headers}')\n",
    "            return headers\n",
    "        \n",
    "            parts = input_filename.split('.')\n",
    " \n",
    "def text_file_record_generator(filename, compression):\n",
    "    \"\"\"\n",
    "    A helper function that returns an object that we can use to iterate through lines in the input file\n",
    "    Args:\n",
    "        filename (str): The input filename\n",
    "        compression (str): The compression type of the input file or None if its not compressed    \n",
    "    \"\"\"\n",
    "    if compression is None: compression = infer_compression(filename)\n",
    "    if compression == None or compression == '':\n",
    "        return open(filename, 'r')\n",
    "    if compression == 'gzip':\n",
    "        import gzip\n",
    "        return gzip.open(filename, 'r')\n",
    "    if compression == 'bz2':\n",
    "        import bz2\n",
    "        return bz2.BZ2File(filename, 'r')\n",
    "    if compression == 'zip':\n",
    "        import zipfile\n",
    "        zf = zipfile.ZipFile(filename, mode, zipfile.ZIP_DEFLATED)\n",
    "        zip_names = zf.namelist()\n",
    "        if len(zip_names) == 0: raise ValueError(f'zero files found in ZIP file {filename}')\n",
    "        if len(zip_names) > 1: raise ValueError(f'{len(zip_names)} files found in ZIP file {filename} {zip_names}.  Only 1 allowed')\n",
    "        return zf.open(zip_names.pop())\n",
    "    if compression == 'xz':\n",
    "        import_lzma\n",
    "        return lzma.LZMAFile(filename, 'r')\n",
    "    raise ValueError(f'unrecognized compression: {compression} for file: {filename}')\n",
    "\n",
    "def base_date_filename_mapper(input_file_path):\n",
    "    \"\"\"\n",
    "    A helper function that parses out the date from a filename.  For example, given a file such as \"/tmp/spx_2018-08-09\", this parses out the \n",
    "    date part of the filename and returns milliseconds (no fractions) since the epoch to that date.\n",
    "    \n",
    "    Args:\n",
    "    \n",
    "        input_filepath (str): Full path to the input file\n",
    "    \n",
    "    Returns:\n",
    "    \n",
    "       int: Milliseconds since unix epoch to the date implied by that file\n",
    "    \n",
    "    >>> base_date_filename_mapper(\"/tmp/spy_1970-1-2_quotes.gz\")\n",
    "    86400000\n",
    "    \"\"\"\n",
    "    filename = os.path.basename(input_file_path)\n",
    "    base_date = dateutil.parser.parse(filename, fuzzy=True)\n",
    "    return round(millis_since_epoch(base_date))\n",
    "\n",
    "def create_text_file_processor(record_generator, line_filter, record_parser, bad_line_handler, record_filter, missing_data_handler,\n",
    "                               aggregators, skip_rows = 1): \n",
    "    return TextFileProcessor(record_generator,\n",
    "                                line_filter,\n",
    "                                record_parser,\n",
    "                                bad_line_handler,\n",
    "                                record_filter,\n",
    "                                missing_data_handler,\n",
    "                                aggregators,\n",
    "                                skip_rows)\n",
    "\n",
    "def get_field_indices(field_names, headers):\n",
    "    \"\"\"\n",
    "    Helper function to get indices of field names in a list of headers\n",
    "    \n",
    "    Args:\n",
    "        field_names (list of str): The fields we want indices of\n",
    "        headers (list of str): All headers\n",
    "        \n",
    "    Returns:\n",
    "        list of int: indices of each field name in the headers list\n",
    "    \"\"\"\n",
    "    field_ids = np.ones(len(field_names), dtype = np.int) * -1\n",
    "    for i, field_name in enumerate(field_names):\n",
    "        if field_name not in headers: raise ParseException(f'{field_name} not in {headers}')\n",
    "        field_ids[i] = headers.index(field_name)\n",
    "    return field_ids\n",
    "\n",
    "def process_marketdata_file(input_filename,\n",
    "                 output_file_prefix_mapper,\n",
    "                 record_parser_creator,\n",
    "                 aggregator_creator,\n",
    "                 line_filter = None, \n",
    "                 compression = None,\n",
    "                 base_date_mapper = None,\n",
    "                 file_processor_creator = create_text_file_processor,\n",
    "                 header_parser_creator = lambda record_generator :  TextHeaderParser(record_generator),\n",
    "                 header_record_generator = text_file_record_generator,\n",
    "                 record_generator = TextFileDecompressor(),\n",
    "                 bad_line_handler = PrintBadLineHandler(),\n",
    "                 record_filter = None,\n",
    "                 missing_data_handler = PriceQtyMissingDataHandler(), \n",
    "                 writer_creator = ArrowWriterCreator()):\n",
    "    \n",
    "    \"\"\"\n",
    "    Processes a single market data file\n",
    "    \n",
    "    Args:\n",
    "        input_filename (str):\n",
    "        output_file_prefix_mapper: A function that takes an input filename and returns the corresponding output filename we want\n",
    "        record_parser_creator:  A function that takes a date and a list of column names and returns a \n",
    "            function that can take a list of fields and return a subclass of Record\n",
    "        line_filter (optional): A function that takes a line and decides whether we want to keep it or discard it.  Defaults to None\n",
    "        compression (str, optional): Compression type for the input file.  Defaults to None\n",
    "        base_date_mapper (optional): A function that takes an input filename and returns the date implied by the filename, \n",
    "            represented as millis since epoch.  Defaults to helper :obj:`function base_date_filename_mapper`\n",
    "        file_processor_creator (optional): A function that returns an object that we can use to iterate through lines in a file.  Defaults to\n",
    "            helper function :obj:`create_text_file_processor`\n",
    "        bad_line_handler (optional): A function that takes a line that we could not parse, and either parses it or does something else\n",
    "            like recording debugging info, or stopping the processing by raising an exception.  Defaults to helper function \n",
    "            :obj:`PrintBadLineHandler`\n",
    "        record_filter (optional): A function that takes a parsed TradeRecord, QuoteRecord, OpenInterestRecord or OtherRecord and decides whether we\n",
    "            want to keep it or discard it.  Defaults to None\n",
    "        missing_data_handler (optional):  A function that takes a parsed TradeRecord, QuoteRecord, OpenInterestRecord or OtherRecord, and decides\n",
    "            deals with any data that is missing in those records.  For example, 0 for bid could be replaced by NAN.  Defaults to helper function:\n",
    "            :obj:`price_qty_missing_data_handler`\n",
    "        writer_creator (optional): A function that takes an output_file_prefix, schema, whether to create a batch id file, and batch_size\n",
    "            and returns a subclass of :obj:`Writer`.  Defaults to helper function: :obj:`arrow_writer_creator`\n",
    "    \"\"\"\n",
    "    \n",
    "    output_file_prefix = output_file_prefix_mapper(input_filename)\n",
    "    \n",
    "    base_date = 0\n",
    "    \n",
    "    if base_date_mapper is not None: base_date = base_date_mapper(input_filename)\n",
    "    \n",
    "    if not is_newer(input_filename, output_file_prefix + '.done'):\n",
    "        print(f'{output_file_prefix + \".done\"} exists and is not older than: {input_filename}')\n",
    "        return\n",
    "                                \n",
    "    header_parser = header_parser_creator(header_record_generator)\n",
    "    print(f'starting file: {input_filename}')\n",
    "    if compression is None: compression = infer_compression(input_filename)\n",
    "    headers = header_parser(input_filename, compression)\n",
    "    \n",
    "    record_parser = record_parser_creator(base_date, headers)\n",
    "    \n",
    "    aggregators = aggregator_creator(writer_creator, output_file_prefix)\n",
    "\n",
    "    file_processor = file_processor_creator(\n",
    "        record_generator, \n",
    "        line_filter, \n",
    "        record_parser, \n",
    "        bad_line_handler, \n",
    "        record_filter, \n",
    "        missing_data_handler,\n",
    "        aggregators\n",
    "    )\n",
    "\n",
    "    start = timer()\n",
    "    if compression is None: compression = \"\"\n",
    "    lines_processed = file_processor(input_filename, compression)\n",
    "    end = timer()\n",
    "    duration = round((end - start) * 1000)\n",
    "    touch(output_file_prefix + '.done')\n",
    "    print(f\"processed: {input_filename} {lines_processed} lines in {duration} milliseconds\")\n",
    "                    \n",
    "def process_marketdata(input_filename_provider, file_processor, num_processes = None, raise_on_error = True):\n",
    "    \"\"\"\n",
    "    Top level function to process a set of market data files\n",
    "    \n",
    "    Args:\n",
    "        input_filename_provider: A function that returns a list of filenames (incl path) we need to process.\n",
    "        file_processor: A function that takes an input filename and processes it, returning number of lines processed. \n",
    "        num_processes (int, optional): The number of processes to run to parse these files.  If set to None, we use the number of cores\n",
    "            present on your machine.  Defaults to None\n",
    "        raise_on_error (bool, optional): If set, we raise an exception when there is a problem with parsing a file, so we can see a stack\n",
    "            trace and diagnose the problem.  If not set, we print the error and continue.  Defaults to True\n",
    "    \"\"\"\n",
    "    \n",
    "    input_filenames = input_filename_provider()\n",
    "    if sys.platform in [\"win32\", \"cygwin\"] and num_processes > 1:\n",
    "        raise Exception(\"num_processes > 1 not supported on windows\")\n",
    "     \n",
    "    if num_processes == 1 or sys.platform in [\"win32\", \"cygwin\"]:\n",
    "        for input_filename in input_filenames:\n",
    "            try:\n",
    "                file_processor(input_filename)\n",
    "            except Exception as e:\n",
    "                new_exc = type(e)(f'Exception: {str(e)}').with_traceback(sys.exc_info()[2])\n",
    "                if raise_on_error: \n",
    "                    raise new_exc\n",
    "                else: \n",
    "                    print(str(new_exc))\n",
    "                    continue\n",
    "    else:\n",
    "        with concurrent.futures.ProcessPoolExecutor(num_processes) as executor:\n",
    "            fut_filename_map = {}\n",
    "            for input_filename in input_filenames:\n",
    "                fut = executor.submit(file_processor, input_filename)\n",
    "                fut_filename_map[fut] = input_filename\n",
    "            for fut in concurrent.futures.as_completed(fut_filename_map):\n",
    "                try:\n",
    "                    fut.result()\n",
    "                    if VERBOSE: print(f'done filename: {fut_filename_map[fut]}')\n",
    "                except Exception as e:\n",
    "                    new_exc = type(e)(f'Exception: {str(e)}').with_traceback(sys.exc_info()[2])\n",
    "                    if raise_on_error: \n",
    "                        raise new_exc\n",
    "                    else: \n",
    "                        print(str(new_exc))\n",
    "                        continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "run_control": {
     "frozen": false,
     "read_only": false
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
