False positive "not-context-manager" for psycopg.connect when used as context manager

[tuukkamustonen]

Bug description

Following:

import psycopg

with psycopg.connect('...') as conn:
    pass

...raises E1129 (not-context-manager) violation.

But the code works fine and sources look ok as psycopg.connect() returns Connection class (which indeed has __enter__ and __exit__).

I'm running psycopg==3.0.1.

Configuration

No response

Command used

pylint

Pylint output

************* Module foo
foo.py:3:0: E1129: Context manager 'NoneType' doesn't implement __enter__ and __exit__. (not-context-manager)

--------------------------------------------------------------------
Your code has been rated at -6.67/10 (previous run: -6.67/10, +0.00)

Expected behavior

There shouldn't be error.

Pylint version

pylint 2.11.1
astroid 2.8.4
Python 3.8.12 (default, Nov  6 2021, 14:26:37) 
[GCC 9.3.0]

OS / Environment

Kubuntu 20.04 LTS

Additional dependencies

No response

[tuukkamustonen]

Possibly related: python/mypy#11004 (something to do with @overloads?)

[rchen152]

I also ran into this issue (with a different contextmanager that also uses overloads). What I discovered is that pylint appears to be checking that the overloads return a contextmanager, which produces an error because overloads return None. That is:

@typing.overload
def some_func():
  ...  # this makes pylint unhappy when it sees `with some_func()` because it looks like `some_func` returns None.

def some_func():
  return MyContextManager()

[rchen152]

A workaround that worked for me is adding return statements to the overloads:

@typing.overload
def some_func():
  return MyContextManager()

[...]

Of course, this is rather ugly and only works if you have control over the code containing the overloads.

[bermed28]

I also faced issues with this, using psycopg[binary] (aka Pscyopg3). I wa using this code per se:

with connect(**connection_params) as conn:
    with conn.cursor() as cursor:
        # some code here
        with cursor.copy(f"COPY {PG_DB_SCHEMA}.{PG_TABLE_NAME} ({columns}) FROM STDIN") as copy:
            # some code here

And when running PyLint I got the error:

************* Module path.to.my.module
         path/to/my/module.py:95:8: E1129: Context manager 'NoneType' doesn't implement __enter__ and __exit__. (not-context-manager)

[adamtuft]

I have come across this issue when using my own context manager with an overloaded factory function on a class. Strangely this doesn't happen when the factory function is free-standing. Here's some demo code:

from enum import Enum, auto
from typing import overload, Literal

class Mode(Enum):
    READ = auto()
    WRITE = auto()
    SIMULATE = auto()

class ConnectionBase:

    def __enter__(self):
        print(f"{type(self).__name__}: enter")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"{type(self).__name__}: exit")
        return exc_type is None

class ReadConnection(ConnectionBase):
    pass

class WriteConnection(ConnectionBase):
    pass

class SimConnection(ConnectionBase):
    pass

def good_code():

    @overload
    def connect(mode: Literal[Mode.READ]) -> ReadConnection: ...

    @overload
    def connect(mode: Literal[Mode.WRITE]) -> WriteConnection: ...

    @overload
    def connect(mode: Literal[Mode.SIMULATE]) -> SimConnection: ...

    def connect(mode: Mode):
        if mode == Mode.READ:
            return ReadConnection()
        if mode == Mode.WRITE:
            return WriteConnection()
        if mode == Mode.SIMULATE:
            return SimConnection()
        raise ValueError(f"unknown connection mode: {mode}")

    con = connect(mode=Mode.READ)
    with con:  # pylint knows that con is a contextmanager
        ...


def bad_code():

    class Project:

        @overload
        def connect(self, mode: Literal[Mode.READ]) -> ReadConnection: ...

        @overload
        def connect(self, mode: Literal[Mode.WRITE]) -> WriteConnection: ...

        @overload
        def connect(self, mode: Literal[Mode.SIMULATE]) -> SimConnection: ...

        def connect(self, mode: Mode):
            if mode == Mode.READ:
                return ReadConnection()
            if mode == Mode.WRITE:
                return WriteConnection()
            if mode == Mode.SIMULATE:
                return SimConnection()
            raise ValueError(f"unknown connection mode: {mode}")

    project = Project()

    con = project.connect(mode=Mode.READ)
    with con: # pylint reports E1129: 'NoneType' doesn't implement __enter__ and __exit__
        ...

Pylint version:

> python3 -m pylint --version
pylint 3.1.0
astroid 3.1.0
Python 3.11.8 (main, Feb 12 2024, 14:50:05) [GCC 13.2.1 20230801]

[adamtuft]

I did a little investigating and wrote a (non-library-specific) test case which fails due to this false-positive. With astroid 3.2.0-dev0 this bug is fixed, thanks to pylint-dev/astroid#1173

[Pierre-Sassoulas]

Sounds great @adamtuft, would you mind adding this test case directly in pylint with a pull request ?

[adamtuft]

Sure thing, I'll put one together 👍