wrong types for csv.DictWriter

[kanetkarster]

Hi, it seems that csv.DictWriter is assuming dictionary keys will always be strings
I've found it useful to use IntEnum keys with DictWriter to avoid having magic strings and write metadata rows positionally

I think it would make sense to change csv.pyi to change fieldnames and the key of _DictRow to be parameterize around a TypeVar('T')

I can add a pull request if you all think this is the right way to go

import csv
import enum
import io

from typing import TextIO

class MyEnum(enum.IntEnum):
    foo = 0
    bar = 1
    baz = 2

def write_data(sink: TextIO) -> None:
    writer = csv.DictWriter(sink, fieldnames=list(MyEnum))
    writer.writerow({0: "<timestamp>", 1: "metadata"})
    writer.writerow({MyEnum.foo: 1000, MyEnum.bar: 200, MyEnum.baz: 300})

def test_write_data() -> None:
    sio = io.StringIO()
    write_data(sio)

    sio.seek(0)
    assert sio.readlines() == [
        "<timestamp>,metadata,\r\n",
        "1000,200,300\r\n",
    ]

mypy ouput:

$ mypy --strict csv_writer.py
csv_writer.py:16: error: Dict entry 0 has incompatible type "int": "str"; expected "str": "Any"
csv_writer.py:16: error: Dict entry 1 has incompatible type "int": "str"; expected "str": "Any"
csv_writer.py:17: error: Dict entry 0 has incompatible type "MyEnum": "int"; expected "str": "Any"
csv_writer.py:17: error: Dict entry 1 has incompatible type "MyEnum": "int"; expected "str": "Any"
csv_writer.py:17: error: Dict entry 2 has incompatible type "MyEnum": "int"; expected "str": "Any"
Found 5 errors in 1 file (checked 1 source file)

This is on Python 3.8.0 with mypy 0.790

[Akuli]

The dict keys can be any type compatible with _csv._writer.writerow, which is currently Any-typed in the stubs. But at runtime, it doesn't work with e.g. enum.Enum, even though it works with at least int, float and str (and also with IntEnum because it inherits from int). I guess you could use a typevar like this:

T = TypeVar("T", float, str)

This typevar can be for example str, float, int (because mypy thinks that int inherits from float even though it actually doesn't) and IntEnum (because it inherits from int). But it can't be enum.Enum.

I guess DictWriter should be made generic. Then your code would type-check as is because DictWriter(...) defaults to DictWriter[Any], and if you want better type checking, you could do DictWriter[MyEnum].

[kanetkarster]

Awesome! I will take a stab at implementing this

[Akuli]

Actually, it accepts any type, because the underlying _csv.writer.writerow converts everything to string:

>>> import csv, enum, sys
>>> class Foo(enum.Enum):
...     X=1
...     Y=2
... 
>>> w=csv.DictWriter(sys.stdout, list(Foo))
>>> w.writeheader()
Foo.X,Foo.Y
>>> w.writerow({Foo.X: "lol", Foo.Y: "wat"})
lol,wat
9

In your case, you needed IntEnum to ensure that MyEnum.foo == 0. It needs that for finding the key 0 from the fieldnames list.