Search code examples
pythonmypy

type annotations for csv.DictWriter


I have trouble typo annotation csv.DictWriter. I am using mypy 0.910 and python 3.9.9.

def csv_bytes_from_dict_list(dict_list: List[Dict[str, Union[str, int]]]) -> bytes:
    with TextIOWrapper(BytesIO(), encoding="utf-8") as r:
        w = DictWriter(r, fieldnames=dict_list[0].keys(), quotechar='"', quoting=QUOTE_ALL)
        w.writeheader()
        w.writerows(dict_list)
        r.flush()
        r.buffer.seek(0)
        b = r.buffer.read()
    return b

Mypy is showing the below errors. I am not sure how to fix them.

test.py: note: In function "csv_bytes_from_dict_list":
test.py:10:13: error: Need type annotation for "w"  [var-annotated]
            w = DictWriter(r, fieldnames=dict_list[0].keys(), quotechar='"', quoting=QUOTE_ALL)
                ^
test.py:10:38: error: Argument "fieldnames" to "DictWriter" has incompatible type "KeysView[str]"; expected "Sequence[<nothing>]"  [arg-type]
            w = DictWriter(r, fieldnames=dict_list[0].keys(), quotechar='"', quoting=QUOTE_ALL)

When I do

w: DictWriter = ...

I get the following error

error: Missing type parameters for generic type "DictWriter"

When I do:

w: DictWriter[Dict[str, Union[str, int]]] = ...

I get the following error.

Argument 1 to "writerows" of "DictWriter" has incompatible type "List[Dict[str, Union[str, int]]]"; expected "Iterable[Mapping[Dict[str, Union[str, int]], Any]]" 

As for the list, the following doesn't help.

fieldnames=list(dict_list[0])

It shows

Argument 1 to "list" has incompatible type "Dict[str, Union[str, int]]"; expected "Iterable[Dict[str, Union[str, int]]]"  [arg-type]
                r, fieldnames=list(dict_list[0]), quotechar='"', quoting=QUOTE_ALL

Solution

  • So by playing around with the typing of the DictWriter, I noticed that whatever type I stick in there will show as the expected type for the first type of the write rows method, essentially the key of the mapping.

    For example

    # if w is typed like this
    w: DictWriter[Dict[str, Union[str, int]]
    # the arguments for w.writerows expects this type
    Iterable[Mapping[Dict[str, Union[str, int]], Any]]
    
    # if the type for w is
    w: DictWriter[dict]
    # the expected type for w.writerows is 
    Iterable[Mapping[Dict[Any, Any], Any]]
    

    Based on that, I came to the conclusion that the type for the DictWriter should be the type of the keys in the mapping.

    # only the type of the keys
    w: DictWriter[str]
    # now the expected type for w.writerows also makes sense
    # althogh its showing 'Any' as second type. I am not sure how to impove that
    Iterable[Mapping[str, Any]]
    

    That combined with list(dict_list[0]) solves the typing errors.

    w: DictWriter[str] = DictWriter(r, fieldnames=list(dict_list[0]), ...)