{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%typecheck\n",
    "import glob\n",
    "import os\n",
    "import sys\n",
    "import re\n",
    "from dateutil import parser as dateutil_parser\n",
    "import numpy as np\n",
    "import concurrent\n",
    "import concurrent.futures\n",
    "from timeit import default_timer as timer\n",
    "import multiprocessing\n",
    "from pyqstrat.pq_utils import infer_compression, millis_since_epoch, touch, get_child_logger\n",
    "from pyqstrat.pyqstrat_cpp import TextFileDecompressor, TextFileProcessor, PrintBadLineHandler, PriceQtyMissingDataHandler\n",
    "from pyqstrat.pyqstrat_cpp import HDF5WriterCreator, Aggregator, Schema, Record, Writer\n",
    "from typing import Sequence, Optional, Callable, Iterable, Union, Tuple\n",
    "\n",
    "\n",
    "RecordGeneratorType = Union[Iterable[str], Iterable[bytes]]\n",
    "RecordGeneratorCreatorType = Callable[[str, str], RecordGeneratorType]\n",
    "RecordParserType = Callable[[Sequence[str]], Record]\n",
    "RecordParserCreatorType = Callable[[int, Sequence[str]], RecordParserType]\n",
    "HeaderParserType = Callable[[str, str], Sequence[str]]\n",
    "HeaderParserCreatorType = Callable[[RecordGeneratorCreatorType], HeaderParserType]\n",
    "LineFilterType = Callable[[str], bool]\n",
    "RecordFilterType = Callable[[Record], bool]\n",
    "BadLineHandlerType = Callable[[str, Exception], Record]\n",
    "MissingDataHandlerType = Callable[[Record], None]\n",
    "FileProcessorType = Callable[[str, Optional[str]], int]\n",
    "InputFileNameProviderType = Callable[[], Sequence[Tuple[str, str]]]\n",
    "WriterCreatorType = Callable[[str, Schema, bool, int], Writer]\n",
    "OutputFilePrefixMapperType = Callable[[str], str]\n",
    "BaseDateMapperType = Callable[[str], int]\n",
    "AggregatorCreatorType = Callable[[WriterCreatorType], Sequence[Aggregator]]\n",
    "FileProcessorCreatorType = Callable[[\n",
    "    RecordGeneratorType, \n",
    "    Optional[LineFilterType], \n",
    "    RecordParserType, \n",
    "    BadLineHandlerType, \n",
    "    Optional[RecordFilterType], \n",
    "    MissingDataHandlerType, \n",
    "    Sequence[Aggregator]],\n",
    "    FileProcessorType]\n",
    "\n",
    "_logger = get_child_logger('pyqstrat')\n",
    "\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: str, include_pattern: str = None, exclude_pattern: str = None) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            path: A pattern such as \"/tmp/abc*.gz\"\n",
    "            include_pattern: Given a pattern such as \"xzy\", will return only filenames that contain xyz\n",
    "            exclude_pattern: 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) -> Sequence[Tuple[str, str]]:\n",
    "        '''\n",
    "        Returns:\n",
    "            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",
    "        compression_tmp = [infer_compression(filename) for filename in files]\n",
    "        compression: Sequence[str] = list(map(lambda x: '' if x is None else x, compression_tmp))\n",
    "        \n",
    "        return list(zip(files, compression))\n",
    "    \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: str) -> None:\n",
    "        '''\n",
    "        Args:\n",
    "            output_dir: 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: str) -> str:\n",
    "        '''\n",
    "        Args:\n",
    "            input_filepath: The input file that we are creating an output file for, e.g. \"/home/xzy.gz\"\n",
    "        Returns:\n",
    "            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",
    "            _logger.debug(f'got input file: {input_filename}')\n",
    "        output_prefix = os.path.join(dirname, input_filename)\n",
    "        return output_prefix\n",
    "\n",
    "\n",
    "class TextHeaderParser:\n",
    "    '''\n",
    "    Parses column headers from a text file containing market data\n",
    "    '''\n",
    "    def __init__(self, \n",
    "                 record_generator_creator: RecordGeneratorCreatorType, \n",
    "                 skip_rows: int = 0, \n",
    "                 separator: str = ',', \n",
    "                 make_lowercase: bool = True) -> None:\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: Number of rows to skip before starting to read the file.  Default is 0\n",
    "            separator: Separator for headers.  Defaults to ,\n",
    "            make_lowercase: Whether to convert headers to lowercase before returning them\n",
    "        '''\n",
    "        self.record_generator_creator = record_generator_creator\n",
    "        self.skip_rows = 0\n",
    "        self.separator = separator\n",
    "        self.make_lowercase = make_lowercase\n",
    "        \n",
    "    def __call__(self, input_filename: str, compression: str) -> Sequence[str]:\n",
    "        '''\n",
    "        Args:\n",
    "        \n",
    "        input_filename The file to read\n",
    "        compression: Compression type, e.g. \"gzip\", or None if the file is not compressed\n",
    "        \n",
    "        Returns:\n",
    "            Column headers\n",
    "        '''\n",
    "        decode_needed = (compression is not None and compression != \"\")\n",
    "        \n",
    "        f = self.record_generator_creator(input_filename, compression)\n",
    "        headers = None\n",
    "        for line_num, line in enumerate(f):\n",
    "            if decode_needed: line = line.decode()  # type: ignore # str does not have decode\n",
    "            if line_num < self.skip_rows: continue\n",
    "            headers = line.split(self.separator)  # type: ignore  # byte does not have split\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",
    "        _logger.debug(f'Found headers: {headers}')\n",
    "        return headers            \n",
    "            \n",
    "\n",
    "def text_file_record_generator_creator(filename: str, compression: str = None) -> RecordGeneratorType:\n",
    "    '''\n",
    "    A helper function that returns a generator that we can use to iterate through lines in the input file\n",
    "    Args:\n",
    "        filename: The input filename\n",
    "        compression: 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 is 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='r', compression=zipfile.ZIP_DEFLATED)\n",
    "        zip_infos = zf.infolist()\n",
    "        zip_names = [zi.filename for zi in zip_infos if not zi.is_dir()]\n",
    "        if len(zip_names) == 0: raise ValueError(f'zero files found in ZIP file {filename}')\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",
    "\n",
    "def base_date_filename_mapper(input_file_path: str) -> int:\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",
    "        input_filepath (str): Full path to the input file\n",
    "    \n",
    "    Returns:\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",
    "\n",
    "def create_text_file_processor(\n",
    "        record_generator: RecordGeneratorType, \n",
    "        line_filter: Optional[LineFilterType], \n",
    "        record_parser: RecordParserType,\n",
    "        bad_line_handler: BadLineHandlerType, \n",
    "        record_filter: Optional[RecordFilterType],\n",
    "        missing_data_handler: MissingDataHandlerType,\n",
    "        aggregators: Sequence[Aggregator],\n",
    "        skip_rows: int = 1) -> FileProcessorType:\n",
    "    \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",
    "\n",
    "def get_field_indices(field_names: Sequence[str], headers: Sequence[str]) -> Sequence[int]:\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=int) * -1\n",
    "    for i, field_name in enumerate(field_names):\n",
    "        if field_name not in headers: raise Exception(f'{field_name} not in {headers}')\n",
    "        field_ids[i] = headers.index(field_name)\n",
    "    return list(field_ids)\n",
    "\n",
    "\n",
    "def process_marketdata_file(\n",
    "        input_filename: str,\n",
    "        compression: str,\n",
    "        output_file_prefix_mapper: OutputFilePrefixMapperType,\n",
    "        record_parser_creator: RecordParserCreatorType,\n",
    "        aggregator_creator: AggregatorCreatorType,\n",
    "        line_filter: LineFilterType = None, \n",
    "        base_date_mapper: BaseDateMapperType = None,\n",
    "        file_processor_creator: FileProcessorCreatorType = create_text_file_processor,\n",
    "        header_parser_creator: HeaderParserCreatorType = lambda record_generator_creator: TextHeaderParser(record_generator_creator),\n",
    "        header_record_generator: RecordGeneratorCreatorType = text_file_record_generator_creator,\n",
    "        record_generator: RecordGeneratorType = TextFileDecompressor(),\n",
    "        bad_line_handler: BadLineHandlerType = PrintBadLineHandler(),\n",
    "        record_filter: RecordFilterType = None,\n",
    "        missing_data_handler: MissingDataHandlerType = PriceQtyMissingDataHandler(), \n",
    "        writer_creator: WriterCreatorType = None) -> None:\n",
    "    '''\n",
    "    Processes a single market data file\n",
    "    \n",
    "    Args:\n",
    "        input_filename: File name (including path) to process\n",
    "        compression: Compression type for the input file.  If not set, we try to infer the compression type from the filename.\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",
    "        aggregator_creator: A function that takes a writer creator and returns a list of Aggregators\n",
    "        line_filter: A function that takes a line and decides whether we want to keep it or discard it.  Defaults to None\n",
    "        base_date_mapper: A function that takes an input filename and returns the date implied by the filename, \n",
    "            represented as millis since epoch.\n",
    "        file_processor_creator: 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",
    "        header_record_generator: A function that takes a filename and compression and returns a generator that we can use to get column headers\n",
    "        record_generator: A function that takes a filename and compression and returns a generator that we \n",
    "            can use to iterate through lines in the file\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 :obj:`HDF5WriterCreatorr`\n",
    "    '''\n",
    "    \n",
    "    if writer_creator is None: writer_creator = HDF5WriterCreator(output_file_prefix_mapper(input_filename), ' ')\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",
    "    header_parser = header_parser_creator(header_record_generator)\n",
    "    _logger.info(f'starting file: {input_filename}')\n",
    "    if compression == \"\": \n",
    "        compression_tmp = infer_compression(input_filename)\n",
    "        compression = '' if compression_tmp is None else compression_tmp  # In C++ don't want virtual function with default argument, so don't allow it here\n",
    "\n",
    "    headers = header_parser(input_filename, compression)  # type: ignore # cannot be None at this point.  \n",
    "    \n",
    "    record_parser = record_parser_creator(base_date, headers)\n",
    "    \n",
    "    aggregators = aggregator_creator(writer_creator)\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",
    "    lines_processed = file_processor(input_filename, compression)\n",
    "    end = timer()\n",
    "    duration = round((end - start) * 1000)\n",
    "    touch(output_file_prefix + '.done')\n",
    "    _logger.info(f\"processed: {input_filename} {lines_processed} lines in {duration} milliseconds\")\n",
    "                    \n",
    "\n",
    "def process_marketdata(\n",
    "        input_filename_provider: InputFileNameProviderType,\n",
    "        file_processor: FileProcessorType,\n",
    "        num_processes: int = None,\n",
    "        raise_on_error: bool = True) -> None:\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 is not None and num_processes > 1:\n",
    "        raise Exception(\"num_processes > 1 not supported on windows\")\n",
    "     \n",
    "    if num_processes is None: num_processes = multiprocessing.cpu_count()\n",
    "        \n",
    "    if num_processes == 1 or sys.platform in [\"win32\", \"cygwin\"]:\n",
    "        for (input_filename, compression) in input_filenames:\n",
    "            try:\n",
    "                file_processor(input_filename, compression)\n",
    "                _logger.debug(f'done: {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, compression) in input_filenames:\n",
    "                fut = executor.submit(file_processor, input_filename, compression)\n",
    "                fut_filename_map[fut] = input_filename\n",
    "            for fut in concurrent.futures.as_completed(fut_filename_map):\n",
    "                try:\n",
    "                    fut.result()\n",
    "                    _logger.debug(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"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
