Anotate return type with psycopg2 type stub

Question

I have a function which returns a psycopg2 connection, if a connection can be established. So the return type should be Optional[psycopg2.connection], or psycopg2.connection | None. However I am unable to import psycopg2.connection at runtime. I've tried the workaround mentioned in How can I import type-definitions from a typeshed stub-file? but that gives me this mypy error: Single overload definition, multiple required. Here's my code

import psycopg2
from typing import Optional, TYPE_CHECKING, overload

if TYPE_CHECKING:
    from psycopg2 import connection
    
    @overload
    def get_connection() -> Optional[connection]: ...


# Make DB error logging less spammy
has_logged_error = False

def get_connection():
    try:
        conn = psycopg2.connect(
            dbname=settings.db_name,
            user=settings.db_user,
            password=settings.db_password,
            host=settings.db_host,
            port=settings.db_port,
        )
        return conn
    except Exception as e:
        global has_logged_error
        if not has_logged_error:
            logger.error(f"Error connecting to DB: {e}")
            has_logged_error = True
        return

Ratiocinate

I cant really reproduce your error at the moment, bc my mypy does not find the psycopg2 type definitions. However, my IDE did not complain, when I removed the overload and put the return type on the actual function definition. – Remarkable

Right, you need to do pip install types-psycopg2. Mypy doesn't complain for me in that case either, but then I get the runtime error that "pysycopg2 has no attribute connection" or some such – Ratiocinate

Sorry, I can't really figure out how to help you here. Maybe you could open an issue in the mypy github repo? – Remarkable

Answer 1

The question you linked proposes some extremely dirty hack which doesn't seem to work any more. There is absolutely no need for it under such simple circumstances. Moreover, to be honest, I cannot reproduce that solution on any mypy version starting from 0.800 (old enough, given that the linked answer is recent), so that perhaps never worked.

I reduced your code samples to contain only minimal return for the sake of readability.

Variant 1: use conditional import and string annotation

import psycopg2
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from psycopg2 import connection
    
def get_connection() -> Optional['connection']:
    return psycopg2.connect(...)

This is simple: mypy known what connection is (defined in stubs); runtime does not try to learn something about connection because it sees simply a string.

Variant 2: use conditional import and annotations future

from __future__ import annotations
import psycopg2
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from psycopg2 import connection
    
def get_connection() -> Optional[connection]:
    return psycopg2.connect(...)

Docs for future imports. This is very similar to direct string usage, but looks nicer and is more convenient, IMO.

Variant 3: use string annotation, but avoid conditional import

import psycopg2
from typing import Optional
    
def get_connection() -> Optional['psycopg2.connection']:
    return psycopg2.connect(...)

Variant 4: use annotations future, but avoid conditional import

from __future__ import annotations
import psycopg2
from typing import Optional
    
def get_connection() -> Optional[psycopg2.connection]:
    return psycopg2.connect(...)

Variants 3 and 4 do not expose that connection is stub-only, hiding it as implementation detail. You may prefer to state that explicitly - then use 1 or 2.

Modification to use current features

This is my favorite. Union syntax is valid in python 3.10+, so if you use an older one - you may want to stick with Optional as described above for consistency. However, annotations future-import makes this expression effectively a string, so if your tools do not perform any runtime typing introspection - you can still use pipe union syntax on older versions. Just be aware that typing.get_type_hints and similar utilities will fail with this syntax on python before 3.10.

from __future__ import annotations
import psycopg2
    
def get_connection() -> psycopg2.connection | None:
    return psycopg2.connect(...)

Tetravalent