Source code for schrodinger.utils.csv_unicode

"""
A wrapper around Python's csv module that allows it to work with Unicode
characters and optionally, gzip compression. The interface is intended
to mirror the csv interface.  See Python's csv documentation for
an explanation of methods and method arguments.
"""

import csv
import gzip
from contextlib import contextmanager

writer = csv.writer


@contextmanager
def reader_open(filename):
    handle = _get_unicode_file_handle(filename, 'rt')
    try:
        yield handle
    finally:
        handle.close()


@contextmanager
def writer_open(filename):
    handle = _get_unicode_file_handle(filename, 'wt')
    try:
        yield handle
    finally:
        handle.close()


def _get_unicode_file_handle(filename, mode):
    """
    Get a file handle with unicode encoding. csv requires the
    handle to be opened in text mode.
    """
    opts = {'mode': mode, 'encoding': "utf-8", 'newline': ""}
    fname = str(filename)  # filename may be a pathlib obj
    if fname.lower().endswith('gz'):
        return gzip.open(filename, **opts)
    return open(filename, **opts)