sqlalchemy-stubs icon indicating copy to clipboard operation
sqlalchemy-stubs copied to clipboard

RowProxy.__getitem__() should accept ColumnElement objects, not just Column

Open sm-Fifteen opened this issue 5 years ago • 0 comments

RowProxy.__getitem__(), works by looking up elements either on its own _keymap or on _parent.keys. In the case of a select statement, the resulting RowProxy object will have keys matching the ColumnElements that were passed to that select, not the initial Column objects, so the signature for __getitem__ should read:

def __getitem__(self, key: Union[str, int, ColumnElement]) -> Any: ...

...instead of:

https://github.com/dropbox/sqlalchemy-stubs/blob/1c36a4ef0898c9bbd0824c0cd2ae8e6b8542930f/sqlalchemy-stubs/engine/result.pyi#L19


Working example:

from sqlalchemy import create_engine, MetaData, Table, Column, select, INTEGER, TEXT

engine = create_engine("sqlite://", echo=True)
test_metadata = MetaData()

test_column = Column('test_id', INTEGER, primary_key=True)
test_table = Table('test_table', test_metadata, test_column)

test_metadata.create_all(engine)

engine.execute(test_table.insert(), [{'test_id': num} for num in range(10)])

res1 = engine.execute(select([test_column]))

for row in res1:
    print("BY LABEL")
    print('\t', row['test_id'])
    print("BY COLUMN")
    print('\t', row[test_column])

test_cast = test_column.cast(TEXT) # Cast[int]/ColumnElement

res2 = engine.execute(select([test_cast]))

for row in res2:
    # The returned label is 'anon_1', not 'test_id'
    # row[test_column] does not exist
    print("BY COLUMN ELEMENT")
    print('\t', row[test_cast]) # Invalid index type "Cast[int]" for "RowProxy"

sm-Fifteen avatar Jan 10 '20 17:01 sm-Fifteen