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: https://github.com/python/mypy/issues/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 https://github.com/pylint-dev/astroid/pull/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 👍

[adamtuft]

Just realised that this is the same issue as https://github.com/pylint-dev/pylint/issues/4660 for which you already have a regression test, fixed by https://github.com/pylint-dev/astroid/pull/1173, so I guess there's no need for me to add another test?

[Pierre-Sassoulas]

Thanks for checking and sorry for the delay. A regression test in pylint would still be useful (especially if we need to upgrade astroid, it's a way to close this pylint issue only when it's actually fixed in pylint because we upgraded astroid and the test that failed now pass). But of course it's less useful than if we had nothing in astroid, I'll let you judge if you still want to contribute :)

[jacobtylerwalls]

@Pierre-Sassoulas I think @adamtuft is pointing out that we actually have this exact case in the test suite, just asserting the wrong result with a sad comment around it. When we upgrade to astroid 3.2 we can just change the result and remove the sad comment 👍.

https://github.com/pylint-dev/pylint/blob/a54d8a68b62d173c872fc1528e16e57453548e2d/tests/functional/r/regression_02/regression_4660.py#L26-L45

Thanks @adamtuft for investigating.

Closing as fixed pending the upgrade to astroid 3.2.