src/solomon/units.py

"""Solomon v1.1 shared sentence/list-item splitter.

units(text)             -> [(start, end), ...]  code-point offsets, sorted, non-overlapping,
                           each text[start:end] is stripped of outer whitespace and non-empty.
resolve(text, quote)    -> (start, end) | None   whitespace-normalised exact match (first hit).
resolve_all(text, quote)-> [(start, end), ...]   every non-overlapping hit.
quote_units(text, quote)-> [unit indices overlapping the resolved quote] ([] if unresolved).

Behaviour builds on solomon.retrieval.candidates (split after .!? + whitespace, and at blank
lines) and adds: hard-wrapped lines are joined (a single newline is NOT a boundary unless the
next line is a bullet / numbered item / "Label:" line, or the current line is a heading-like
short line), abbreviations and initials do not end sentences, and over-long punctuation-free
segments are split at line breaks. The splitter never normalises or drops source characters
other than outer whitespace of a unit.
"""
import re

__all__ = ['units', 'unit_texts', 'resolve', 'resolve_all', 'quote_units', 'span_units',
           'normalise']

ABBREV = {
    'mr', 'mrs', 'ms', 'dr', 'prof', 'sr', 'jr', 'st', 'mt', 'rev', 'hon', 'gen', 'col', 'capt',
    'lt', 'sgt', 'gov', 'sen', 'rep', 'pres', 'supt', 'insp',
    'inc', 'ltd', 'co', 'corp', 'llc', 'plc', 'pty', 'bros', 'dept', 'univ', 'assn', 'assoc',
    'no', 'nos', 'vol', 'vols', 'p', 'pp', 'para', 'paras', 'sec', 'secs', 'art', 'arts', 'ch',
    'cl', 'fig', 'figs', 'ref', 'reg', 'regs', 'sch', 'pt', 'div', 'subdiv', 'r', 's', 'ss',
    'e.g', 'i.e', 'cf', 'vs', 'v', 'viz', 'approx', 'est', 'al', 'etc', 'ibid', 'op', 'resp',
    'jan', 'feb', 'mar', 'apr', 'jun', 'jul', 'aug', 'sep', 'sept', 'oct', 'nov', 'dec',
    'mon', 'tue', 'tues', 'wed', 'thu', 'thur', 'thurs', 'fri', 'sat', 'sun',
    'u.s', 'u.k', 'u.s.a', 'n.a', 'a.m', 'p.m', 'ph.d', 'b.a', 'm.d', 'd.c', 'n.y', 'u.n',
    'ave', 'blvd', 'rd', 'hwy', 'ext', 'tel', 'ph', 'mob', 'attn', 'encl', 'cc', 're',
    'min', 'max', 'mg', 'ml', 'kg', 'hr', 'hrs', 'wk', 'yr', 'yrs', 'dist', 'cir', 'cal',
    'app', 'amend', 'stat', 'ord', 'rul', 'ed', 'eds', 'trans', 'dir', 'mgr', 'asst', 'admin',
}
# Abbreviations that commonly END sentences too; only treat as non-terminal when the next
# word starts lower-case or with a digit.
SOFT_ABBREV = {'etc', 'inc', 'ltd', 'co', 'corp', 'llc', 'plc', 'pty', 'al', 'u.s', 'u.k',
               'a.m', 'p.m', 'jr', 'sr', 'est', 'resp', 'ibid', 'min', 'max', 'hr', 'hrs',
               'mg', 'ml', 'kg', 'mon', 'tue', 'wed', 'thu', 'fri', 'sat', 'sun', 'jan', 'feb',
               'mar', 'apr', 'jun', 'jul', 'aug', 'sep', 'sept', 'oct', 'nov', 'dec', 'd.c',
               'n.y', 'ph.d', 'm.d'}

_BULLET = re.compile(
    r'[ \t]*(?:[-*•●◦‣⁃∙·▪■–—>]\s'
    r'|\(?(?:\d{1,3}(?:\.\d{1,3})*|[a-zA-Z]|[ivxlcdmIVXLCDM]{1,6})[.)]\s'
    r'|\((?:\d{1,3}|[a-zA-Z]|[ivxlcdm]{1,6})\)\s'
    r'|(?:Q|A|Question|Answer)\s*[:.]\s'
    r'|[A-Z][A-Za-z0-9 .,&/\'()\-]{0,38}:[ \t])')
_SENT_END = re.compile(r'[.!?…]+["\'’”)\]]*(?=\s)')
_WS = re.compile(r'\s+')
_WORD_BEFORE = re.compile(r'([A-Za-z][A-Za-z.]*)\.$')
LONG_UNIT = 900
_CO_SUFFIX = {'ltd', 'inc', 'llc', 'plc', 'co', 'corp', 'limited', 'llp', 'lp'}
_ENUM_ONLY = re.compile(r'\(?(?:\d{1,3}(?:\.\d{1,3})*|[a-zA-Z]|[ivxlcdm]{1,6})[.)]')


def _lines(text, a, b):
    """(start, end) of each line in text[a:b], end excluding the newline."""
    out = []
    pos = a
    while pos < b:
        nl = text.find('\n', pos, b)
        end = b if nl < 0 else nl
        out.append((pos, end))
        pos = end + 1
    return out


def _paragraphs(text):
    """Split at blank lines (a line containing only whitespace)."""
    out = []
    start = 0
    for m in re.finditer(r'\n[ \t\r\f\v]*\n\s*', text):
        out.append((start, m.start()))
        start = m.end()
    out.append((start, len(text)))
    return out


def _heading_like(line, maxlen):
    s = line.strip()
    if not s:
        return False
    if s[-1] in '.,;:!?"”\'' and not s.endswith(':'):
        return False
    if s.endswith(':'):
        return True
    words = s.split()
    if len(s) <= 60 and (s.isupper() or maxlen >= 50 and len(s) < 0.6 * maxlen):
        return True
    return len(words) <= 8 and sum(w[:1].isupper() or not w[:1].isalpha() for w in words) == len(words) \
        and len(s) < 0.6 * max(maxlen, 1)


def _blocks(text, a, b):
    """Break a paragraph into blocks at list/label lines and heading-like lines."""
    lines = [(s, e) for s, e in _lines(text, a, b)]
    if not lines:
        return []
    maxlen = max(len(text[s:e].strip()) for s, e in lines)
    wrapped = maxlen >= 50 and len(lines) >= 3
    blocks = []
    cur = lines[0][0]
    for i in range(1, len(lines)):
        ps, pe = lines[i - 1]
        s, e = lines[i]
        prev, line = text[ps:pe], text[s:e]
        brk = bool(_BULLET.match(line))
        if not brk:
            pst, lst = prev.strip(), line.strip()
            cont = lst[:1].islower() and not pst.endswith(':')
            if not wrapped:
                # Unwrapped text: every line is its own block (lists, addresses, form
                # fields) unless the line is an indented or lower-case continuation.
                indented = (len(line) - len(line.lstrip(' \t'))) > (len(prev) - len(prev.lstrip(' \t')))
                brk = not (cont or (indented and pst and pst[-1] not in '.!?:'))
            elif pst and ' ' not in pst and len(pst) > 20 and not cont:
                brk = True  # URL / token-only line
            elif _heading_like(prev, maxlen) and not cont:
                brk = True
        if brk:
            blocks.append((cur, pe))
            cur = s
    blocks.append((cur, lines[-1][1]))
    return blocks


def _is_abbrev(text, dot_end, match_end):
    """dot_end = index just past the terminal punctuation run's first char."""
    token = text[max(0, dot_end - 12):dot_end]
    m = _WORD_BEFORE.search(token)
    if not m:
        return False
    word = m.group(1)
    # Must be a whole word (not the tail of a longer alnum run).
    wstart = dot_end - 1 - len(word)
    if wstart > 0 and text[wstart - 1].isalnum():
        return False
    low = word.lower().rstrip('.')
    nxt = text[match_end:match_end + 40].lstrip()
    nxt_c = nxt[:1]
    if len(word) == 1 and word.isalpha():
        return True  # initial: "J. Smith", "s. 12", "r. 4"
    if '.' in word and len(word) <= 6 and all(len(p) <= 2 for p in word.split('.') if p):
        # dotted acronyms like U.S., e.g., i.e., a.m.
        if low in SOFT_ABBREV:
            return not nxt_c.isupper()
        return True
    if low in ABBREV:
        if low in SOFT_ABBREV:
            nxt_w = re.match(r'[A-Za-z]+', nxt)
            if nxt_w and nxt_w.group(0).lower() in _CO_SUFFIX:
                return True  # "Pty. Ltd.", "Co. Inc."
            return bool(nxt_c) and (nxt_c.islower() or nxt_c.isdigit())
        return True
    return False


def _sentences(text, a, b):
    out = []
    start = a
    for m in _SENT_END.finditer(text, a, b):
        end = m.end()
        # next non-space char
        k = end
        while k < b and text[k].isspace():
            k += 1
        if k >= b:
            continue
        nxt = text[k]
        if nxt.islower() and not text[m.start()] in '!?':
            continue  # "approx. five", "cl. 3 applies" - not a sentence start
        if text[m.start()] == '.' and m.group(0)[:1] == '.' and len(m.group(0).rstrip('"\'’”)]')) == 1 \
                and _is_abbrev(text, m.start() + 1, k):
            continue
        # enumerator like "1." or "(a)." right before a new sentence stays attached behind.
        out.append((start, end))
        start = k
    out.append((start, b))
    return out


def _strip(text, s, e):
    while s < e and text[s].isspace():
        s += 1
    while e > s and text[e - 1].isspace():
        e -= 1
    return s, e


def _meaningful(text, s, e):
    return any(c.isalnum() for c in text[s:e])


def _split_long(text, s, e):
    if e - s <= LONG_UNIT:
        return [(s, e)]
    parts = [(ls, le) for ls, le in _lines(text, s, e)]
    if len(parts) > 1:
        return parts
    # single line with no sentence punctuation: split at ';' then whitespace
    out = []
    cur = s
    while e - cur > LONG_UNIT:
        window = text[cur:cur + LONG_UNIT]
        cut = max(window.rfind('; '), window.rfind(', '))
        if cut < LONG_UNIT // 3:
            cut = window.rfind(' ')
        if cut <= 0:
            cut = LONG_UNIT - 1
        out.append((cur, cur + cut + 1))
        cur = cur + cut + 1
    out.append((cur, e))
    return out


def units(text):
    """Sentence / list-item units as sorted, non-overlapping (start, end) code-point ranges."""
    if not isinstance(text, str):
        raise TypeError('text must be str')
    out = []
    for pa, pb in _paragraphs(text):
        for ba, bb in _blocks(text, pa, pb):
            for sa, sb in _sentences(text, ba, bb):
                for xa, xb in _split_long(text, sa, sb):
                    xa, xb = _strip(text, xa, xb)
                    if xb <= xa:
                        continue
                    if not _meaningful(text, xa, xb):
                        continue  # separators / punctuation-only lines are not units
                    out.append((xa, xb))
    # An enumerator-only unit ("4.", "(b)") is merged into the unit that follows it.
    merged = []
    i = 0
    while i < len(out):
        s, e = out[i]
        if i + 1 < len(out) and _ENUM_ONLY.fullmatch(text[s:e]):
            merged.append((s, out[i + 1][1]))
            i += 2
            continue
        merged.append((s, e))
        i += 1
    return merged


def unit_texts(text):
    return [text[s:e] for s, e in units(text)]


# ---- quote resolution ---------------------------------------------------------------------

_CHARMAP = str.maketrans({'‘': "'", '’': "'", '‚': "'", '‛': "'",
                          '“': '"', '”': '"', '„': '"', ' ': ' ',
                          '‐': '-', '‑': '-', '‒': '-', '–': '-',
                          '—': '-', '−': '-', '­': '­'})


def normalise(text):
    """Collapse whitespace runs to one space and map curly quotes/dashes 1:1 to ASCII.
    Returns (normalised_string, index_map) where index_map[i] = source index of char i;
    index_map has one extra trailing entry = len(text)."""
    chars, idx = [], []
    i, n = 0, len(text)
    t = text.translate(_CHARMAP)
    while i < n:
        c = t[i]
        if c.isspace():
            j = i
            while j < n and t[j].isspace():
                j += 1
            chars.append(' ')
            idx.append(i)
            i = j
            continue
        chars.append(c)
        idx.append(i)
        i += 1
    idx.append(n)
    return ''.join(chars), idx


def _norm_quote(quote):
    q, _ = normalise(quote)
    return q.strip()


def resolve_all(text, quote):
    if not isinstance(quote, str):
        return []
    q = _norm_quote(quote)
    if not q:
        return []
    nt, idx = normalise(text)
    hits = []
    pos = nt.find(q)
    while pos >= 0:
        s = idx[pos]
        e = idx[pos + len(q) - 1] + 1
        hits.append((s, e))
        pos = nt.find(q, pos + len(q))
    return hits


def resolve(text, quote):
    hits = resolve_all(text, quote)
    return hits[0] if hits else None


def span_units(text, span, unit_list=None):
    """Indices of units overlapping a (start, end) char range."""
    us = units(text) if unit_list is None else unit_list
    s, e = span
    return [i for i, (a, b) in enumerate(us) if a < e and s < b]


def quote_units(text, quote, unit_list=None):
    span = resolve(text, quote)
    if span is None:
        return []
    return span_units(text, span, unit_list)