potassco / clorm

🗃️ A Python ORM-like interface for the Clingo Answer Set Programming (ASP) reasoner
https://clorm.readthedocs.io
MIT License
52 stars 5 forks source link

Type annotations for query api #108

Closed florianfischer91 closed 2 years ago

florianfischer91 commented 2 years ago

This PR adds type annotations to the query api. In the past when someone uses fb.query(...).join(...).all() the type of the returned value of all() was Any which is not very usefull for a beginner who is not so familiar how the api works and what is returned. This will change with this PR

Let's take a look at the examples below:

from clorm import Predicate, FactBase

class P6(Predicate):
    a: str
    b: int

class P7(Predicate, is_tuple=True):
    a: int
    b: int

p6s = [P6("1", 1), P6("2", 1), P6("3", 2), P6("4", 2)]
p7s = [P7(1, 1), P7(2, 1), P7(3, 2), P7(4, 2)]

fb = FactBase([*p6s, *p7s])

result = fb.query(P6, P7).join(P6.b == P7.b).where(P6.a == "1").all()
reveal_type(result)

result1 = fb.query(P6, P7).join(P6.b == P7.b).select(P6, P6.b).all()
reveal_type(result1)

result2 = fb.query(P6, P7).join(P6.b == P7.b).group_by(P6).all()
reveal_type(result2)

result3 = fb.query(P6, P7).join(P6.b == P7.b).group_by(P6).count()
reveal_type(result3)

result4 = fb.query(P6, P7).join(P6.b == P7.b).group_by(P6.a, P7.b).select(P6, P7.b).all()
reveal_type(result4)

If you run it with mypy example.py, mypy can infer the correct type of the result now

example.py:17: note: Revealed type is "typing.Generator[Tuple[example.P6, example.P7], None, None]"
example.py:20: note: Revealed type is "typing.Generator[Tuple[example.P6, builtins.int], None, None]"
example.py:23: note: Revealed type is "typing.Generator[Tuple[example.P6, typing.Iterator[Tuple[example.P6, example.P7]]], None, None]"
example.py:26: note: Revealed type is "typing.Iterator[Tuple[example.P6, builtins.int]]"
example.py:29: note: Revealed type is "typing.Generator[Tuple[Tuple[builtins.str, builtins.int], typing.Iterator[Tuple[example.P6, builtins.int]]], None, None]"

and also the correct type is shown in VS-Code grafik

I also added a small TestCase for testing the correct type infered by mypy, but it will be skipped if mypy is not installed.

closes #62