mypy annotations for rows_from_file(), run mypy in CI

Refs #289, #279
This commit is contained in:
Simon Willison 2021-06-22 11:04:32 -07:00
commit 7b3fdf0fcd
4 changed files with 26 additions and 18 deletions

View file

@ -26,11 +26,14 @@ jobs:
- name: Install dependencies - name: Install dependencies
run: | run: |
pip install -e '.[test]' pip install -e '.[test]'
pip install mypy
- name: Optionally install numpy - name: Optionally install numpy
if: matrix.numpy == 1 if: matrix.numpy == 1
run: pip install numpy run: pip install numpy
- name: Run tests - name: Run tests
run: | run: |
pytest pytest
- name: run mypy
run: mypy sqlite_utils
- name: Check formatting - name: Check formatting
run: black . --check run: black . --check

View file

@ -1,6 +1,6 @@
import base64 import base64
import click import click
from click_default_group import DefaultGroup from click_default_group import DefaultGroup # type: ignore
from datetime import datetime from datetime import datetime
import hashlib import hashlib
import pathlib import pathlib

View file

@ -11,7 +11,7 @@ import json
import os import os
import pathlib import pathlib
import re import re
from sqlite_fts4 import rank_bm25 from sqlite_fts4 import rank_bm25 # type: ignore
import sys import sys
import textwrap import textwrap
import uuid import uuid
@ -39,14 +39,14 @@ USING\s+(?P<using>\w+) # e.g. USING FTS5
) )
try: try:
import pandas as pd import pandas as pd # type: ignore
except ImportError: except ImportError:
pd = None pd = None
try: try:
import numpy as np import numpy as np # type: ignore
except ImportError: except ImportError:
np = None np = None # type: ignore
Column = namedtuple( Column = namedtuple(
"Column", ("cid", "name", "type", "notnull", "default_value", "is_pk") "Column", ("cid", "name", "type", "notnull", "default_value", "is_pk")

View file

@ -5,17 +5,18 @@ import enum
import io import io
import json import json
import os import os
from typing import Generator from typing import cast, BinaryIO, Iterable, Optional, Tuple, Type
import click import click
try: try:
import pysqlite3 as sqlite3 import pysqlite3 as sqlite3 # type: ignore
import pysqlite3.dbapi2 import pysqlite3.dbapi2 # type: ignore
OperationalError = pysqlite3.dbapi2.OperationalError OperationalError = pysqlite3.dbapi2.OperationalError
except ImportError: except ImportError:
import sqlite3 # https://github.com/python/mypy/issues/1153#issuecomment-253842414
import sqlite3 # type: ignore
OperationalError = sqlite3.OperationalError OperationalError = sqlite3.OperationalError
@ -143,12 +144,11 @@ class RowsFromFileBadJSON(RowsFromFileError):
def rows_from_file( def rows_from_file(
fp, fp: BinaryIO,
format=None, format: Optional[Format] = None,
dialect=None, dialect: Optional[Type[csv.Dialect]] = None,
encoding=None, encoding: Optional[str] = None,
detect_types=False, ) -> Tuple[Iterable[dict], Format]:
) -> Generator[dict, None, None]:
if format == Format.JSON: if format == Format.JSON:
decoded = json.load(fp) decoded = json.load(fp)
if isinstance(decoded, dict): if isinstance(decoded, dict):
@ -159,8 +159,13 @@ def rows_from_file(
elif format == Format.NL: elif format == Format.NL:
return (json.loads(line) for line in fp if line.strip()), Format.NL return (json.loads(line) for line in fp if line.strip()), Format.NL
elif format == Format.CSV: elif format == Format.CSV:
decoded_fp = io.TextIOWrapper(fp, encoding=encoding or "utf-8-sig") use_encoding: str = encoding or "utf-8-sig"
return csv.DictReader(decoded_fp, dialect=dialect), Format.CSV decoded_fp = io.TextIOWrapper(fp, encoding=use_encoding)
if dialect is not None:
reader = csv.DictReader(decoded_fp, dialect=dialect)
else:
reader = csv.DictReader(decoded_fp)
return reader, Format.CSV
elif format == Format.TSV: elif format == Format.TSV:
return ( return (
rows_from_file( rows_from_file(
@ -170,7 +175,7 @@ def rows_from_file(
) )
elif format is None: elif format is None:
# Detect the format, then call this recursively # Detect the format, then call this recursively
buffered = io.BufferedReader(fp, buffer_size=4096) buffered = io.BufferedReader(cast(io.RawIOBase, fp), buffer_size=4096)
first_bytes = buffered.peek(2048).strip() first_bytes = buffered.peek(2048).strip()
if first_bytes.startswith(b"[") or first_bytes.startswith(b"{"): if first_bytes.startswith(b"[") or first_bytes.startswith(b"{"):
# TODO: Detect newline-JSON # TODO: Detect newline-JSON