myw_csv.py

"""
CSV Parser that handles WKB columns safely
Copyright (c) 2025-2026 IQGeo Group Plc. Use subject to conditions at license.txt
"""

import os.path
import sys
from typing import TYPE_CHECKING
import pandas as pd

if TYPE_CHECKING:
    from io import TextIOWrapper

# override these at runtime if you want to change CSV conventions, but don't edit the file.
SEPARATOR_CHARACTER = ","
DELIMITER_CHARACTER = '"'
ESCAPE_CHARACTER = "\\"


def warn(filename, message, line_no=-1):
    """print a warning to stderror about a problem in a CSV file during parsing."""
    line_ref = ""
    if line_no > 0:
        line_ref = f":l{line_no}"
    sys.stderr.write(f"CSV WARNING: {filename}{line_ref} - {message}\n")
    # ENH include these in the main report JSON etc from the other file somehow.


def _split_allowing_quotes(line: str):
    """in CSV files, commas are allowed inside cells that are delimited by quotes. This method
    splits a line allowing cells to be quoted, and then removes any quote chars used as delimiters.
    Quotes are ignored in the middle/end of a normal cell, but have special meaning at the start
    of cells.
    """
    # strip newlines at end of line.
    while line[-1] in {"\n", "\r"}:
        line = line[:-1]

    this_cell = []
    all_cells = []
    in_quote = False
    next_must_be_comma = False
    always_emit = False

    for cursor, char in enumerate(line):
        if next_must_be_comma:
            if char != SEPARATOR_CHARACTER:
                raise ValueError(f"bad quotes at char {cursor} in line {line}.")
            # next cell.
            all_cells.append("".join(this_cell))
            this_cell = []
            # reset "expect comma" flag.
            next_must_be_comma = False
        elif always_emit:
            # when the previous char was a backslash escape.
            this_cell.append(char)
            always_emit = False
        elif char == ESCAPE_CHARACTER:
            # parse, but don't emit, escapes?
            # this_cell.append(c)
            always_emit = True
        elif in_quote:
            if char != DELIMITER_CHARACTER:
                this_cell.append(char)
            else:
                # If the next char exists and is also DELIMITER_CHARACTER, i.e. "" for escaped
                # quote.
                if cursor + 1 < len(line) and line[cursor + 1] == DELIMITER_CHARACTER:
                    # escaped quote, emit only one quote (the next one) to "decode" the escaped
                    # JSON.
                    always_emit = True
                else:
                    # end the quote, do not emit it.
                    in_quote = False
                    next_must_be_comma = True
        elif len(this_cell) == 0 and char == DELIMITER_CHARACTER:
            in_quote = True
            # do not emit the quote.
        elif char == SEPARATOR_CHARACTER:
            # next cell.
            all_cells.append("".join(this_cell))
            this_cell = []
        else:
            this_cell.append(char)

    # write the last cell
    all_cells.append("".join(this_cell))
    return all_cells


def _detect_BOM(file_path) -> str:
    """Attempt to guess encoding from the first few bytes of the file."""
    encoding = None
    with open(file_path, "rb") as bytefile:
        beginning = bytefile.read(4)
        # The order of these if-statements is important
        # otherwise UTF32 LE may be detected as UTF16 LE as well
        # from: https://stackoverflow.com/a/65841914
        if beginning == b"\x00\x00\xfe\xff":
            encoding = "utf_32_be"
        elif beginning == b"\xff\xfe\x00\x00":
            encoding = "utf_32_le"
        elif beginning[0:3] == b"\xef\xbb\xbf":
            # UTF-8 (this encoding eats the BOM for us)
            encoding = "utf_8_sig"
        elif beginning[0:2] == b"\xff\xfe":
            encoding = "utf_16_le"
        elif beginning[0:2] == b"\xfe\xff":
            encoding = "utf_16_be"
        else:
            # No BOM, we fall back to our default encoding, with the "fast alias":
            encoding = "utf-8"
    return encoding


def _read_logical_lines(file_obj: "TextIOWrapper"):
    """Generator that yields complete logical lines from a CSV file, handling newlines
    inside quoted cells. Accumulated physical lines are joined with their original newline
    characters preserved.
    """
    # Check if this accumulated line forms a complete logical row by tracking quote state
    accumulated_line = ""
    in_quote = False
    cursor = 0  # cursor is only reset to 0 when a line is yielded.
    always_emit = False
    for physical_line in file_obj:
        accumulated_line += physical_line
        # Resume handling the string at the last cursor point (0 = totally new
        # string, >0 = combined line.)
        while cursor < len(accumulated_line):
            char = accumulated_line[cursor]
            if always_emit:
                # when the previous char was a backslash escape.
                cursor += 1
                always_emit = False
            elif char == ESCAPE_CHARACTER:
                # first parse phase, do emit escapes, handled in a later pass.
                always_emit = True
                cursor += 1
            elif char == DELIMITER_CHARACTER:
                # Check if this is an escaped quote (followed by another quote)
                if (
                    cursor + 1 < len(accumulated_line)
                    and accumulated_line[cursor + 1] == DELIMITER_CHARACTER
                ):
                    # Escaped quote "", skip both characters
                    cursor += 2
                else:
                    # Regular quote delimiter, toggle quote state
                    in_quote = not in_quote
                    cursor += 1
            else:
                cursor += 1

        # If we're not in a quote state, we have a complete logical row and reset state.
        if not in_quote:
            yield accumulated_line
            accumulated_line = ""
            cursor = 0

    # Yield any remaining accumulated line (shouldn't happen in well-formed CSV)
    if accumulated_line:
        yield accumulated_line


def _read_csv_raw(file_path: str, encoding=None) -> dict:
    """Read a CSV file into a dictionary of lists, no type conversion.
    {
        "col1": ["1", "2", "3", "4"],
        "col2": ["asdf", "bsdf", "csdf", "dsdf"],
    }

    (If encoding is passed, it assumes that is correct and does a small check on first line for
    leaked BOMs. If encoding is not passed, it tries to auto-detect encoding from the BOM,
    defaulting to UTF-8 if none detected.)
    """
    frame = {}
    index_to_key = {}
    if encoding is None:
        encoding = _detect_BOM(file_path)

    with open(file_path, "r", encoding=encoding) as f:
        first_line = True
        for line_no, line in enumerate(_read_logical_lines(f)):
            # Strip any remaining byte order marks
            if first_line and line[0] == "\ufeff":
                # UTF-16 and UTF-32 BOM.
                line = line[1:]
            elif first_line and line[:3] == "\xef\xbb\xbf":
                # UTF-8 BOM. We only get here if user specified `utf-8` on CLI.
                line = line[3:]
            # Line is like 'id,housing,path\n'
            cells = _split_allowing_quotes(line)
            if first_line:
                # Handle duplicate column names by appending a suffix (per pandas behaviour)
                # but also flag a warning.
                column_counts = {}
                for i, cell in enumerate(cells):
                    # Handle duplicate column names by appending a suffix
                    if cell in column_counts:
                        column_counts[cell] += 1
                        unique_name = f"{cell}.{column_counts[cell]}"
                        warn(
                            os.path.basename(file_path),
                            f"Duplicate key {cell} (relabelling)",
                            1,
                        )
                    else:
                        column_counts[cell] = 1
                        unique_name = cell

                    frame[unique_name] = []
                    index_to_key[i] = unique_name
                first_line = False
            else:
                if len(cells) != len(frame):
                    # Print a warning about the file, and pad with NaN just like the pandas parser.
                    warn(
                        os.path.basename(file_path),
                        "Wrong number of cells in file (truncating)",
                        line_no + 1,
                    )
                    if len(cells) > len(frame):
                        # truncate.
                        cells = cells[: len(frame)]
                    else:
                        # pad with NaN.
                        cells = cells + (["NaN"] * (len(frame) - len(cells)))
                for i, cell in enumerate(cells):
                    key = index_to_key[i]
                    frame[key].append(cell)

    return frame


# NOTE: I did have 0 and 1 in the bools category, but sometimes an enum field only has 0s and 1s in
# it for the whole table, and it should be an int. This could be improved from name context?
def _is_bool(s: str) -> bool:
    valid_bools = {"TRUE", "FALSE", "T", "F", "YES", "NO"}
    return s.upper() in valid_bools


def _cast_bool(s: str) -> bool:
    up = s.upper()
    if up in {"T", "YES", "TRUE"}:
        return True
    elif up in {"F", "NO", "FALSE"}:
        return False
    else:
        raise ValueError(f"invalid bool {s}?")


def _is_hex(s: str) -> bool:
    try:
        int(s, 16)
        return True
    except ValueError:
        return False


def _is_float(s: str) -> bool:
    try:
        float(s)
        return True
    except ValueError:
        return False


def _detect_and_convert_type(l: list[str], full_check=False) -> list:
    """
    input is a list of strings raw from a CSV. This method will detemine the type and convert
    all items.

    Possible types are hexdata, integer, float, string, or boolean.

    Note that boolean does not treat a column of 0s and 1s as bools, they need to be t/f,
    true/false, yes/no.

    Also note that hexdata is returned as a string, since this is the behavior of pandas that we
    are emulating.
    """
    # by default, do nothing.
    cast = lambda i: i

    # A couple of defensive checks here for very short files.
    if len(l) == 0:
        return l

    non_empty_cells = [i for i in l if i != ""]
    non_empty_count = len(non_empty_cells)

    if non_empty_count == 0:
        return [None] * len(l)

    test_item = non_empty_cells[0]
    # for speed, we test various values from the list, unless `full_check` passed (for very short
    # lists, just add the whole rest of the list.)
    if not full_check and non_empty_count > 3:
        verify_items = [
            non_empty_cells[1],
            non_empty_cells[non_empty_count // 2],
            non_empty_cells[-1],
        ]
    else:
        verify_items = non_empty_cells[1:]
    done = False

    # we categorise as either: hexdata, integer, float, string, or boolean.
    if _is_bool(test_item):
        new_cast = _cast_bool
        try:
            for i in verify_items:
                new_cast(i)
            cast = new_cast
            done = True
        except ValueError:
            # test string is a coincidence, try the other tests.
            test_item = i
            verify_items.append(test_item)
    if not done and test_item.isdecimal():
        new_cast = int
        try:
            for i in verify_items:
                new_cast(i)
            cast = new_cast
            done = True
        except ValueError:
            # test string is a coincidence, try the other tests.
            test_item = i
            verify_items.append(test_item)
    if not done and _is_hex(test_item):
        # We return hex encoded WKB as a hex string (i.e. unchanged), as that is how pandas would
        # parse it (if it didn't segfault parsing a float).
        test = lambda i: int(i, 16)
        try:
            for i in verify_items:
                test(i)
            done = True
        except ValueError:
            # test string is a coincidence, try the other tests.
            test_item = i
            verify_items.append(test_item)
    if not done and _is_float(test_item):
        new_cast = float
        try:
            for i in verify_items:
                new_cast(i)
            cast = new_cast
            done = True
        except ValueError:
            # test string is a coincidence, try the other tests.
            test_item = i
            verify_items.append(test_item)

    try:
        # Any empty cells must be None, for pd compatibility.
        return [cast(i) if i != "" else None for i in l]
    except ValueError:
        # failed to correctly guess type - fall back to str. Shouldn't happen if full_check passed.
        # still turn empty string into Nones.
        return [i if i != "" else None for i in l]


def read_csv(file_path: str, encoding=None) -> pd.DataFrame:
    """
    Reads a CSV file into memory safely, and does careful type conversions to turn it into a Pandas
    DataFrame.
    """
    dict_frame = _read_csv_raw(file_path, encoding=encoding)

    dict_frame_with_types = {
        k: _detect_and_convert_type(values) for k, values in dict_frame.items()
    }

    return pd.DataFrame(dict_frame_with_types)


def _test_parser(file_path: str, encoding="utf-8"):
    print("reading file into memory:")
    dict_frame = _read_csv_raw(file_path, encoding=encoding)
    print("determining types...")
    dict_frame_with_types = {}
    for k, values in dict_frame.items():
        new_values = _detect_and_convert_type(values)
        print(f"  {k}: {type(new_values[0])}")
        dict_frame_with_types[k] = new_values
    print("loading into pandas...")
    pf = pd.DataFrame(dict_frame_with_types)
    print(f"pandas types: {pf.dtypes!r}")
    print("success!")
    return pf


if __name__ == "__main__":
    import sys

    _test_parser(sys.argv[1])