Add better type deduction:

This part one of many to add better type deduction to Refurb. Previously I only
pulled type out of `NameExpr` nodes because it was convienent, but this left
out a lot of node types, meaning Refurb could not detect them. This new change
allows for extracting the type of an AST node, meaning Refurb will (in the
future) be much better at detecting certain checks.

This is still in the works and has only been applied to one check, but as time
goes on, this will start to be implemented for more checks.

refurb/checks/common.py

@@ -17,6 +17,7 @@
     Expression,
     FloatExpr,
     ForStmt,
+    FuncDef,
     GeneratorExpr,
     IfStmt,
     IndexExpr,
@@ -34,11 +35,14 @@
     StarExpr,
     Statement,
     StrExpr,
+    SymbolTableNode,
     TupleExpr,
+    TypeAlias,
     TypeInfo,
     UnaryExpr,
+    Var,
 )
-from mypy.types import AnyType, Instance, TupleType, Type
+from mypy.types import AnyType, CallableType, Instance, TupleType, Type
 
 from refurb.error import Error
 from refurb.visitor import TraverserVisitor
@@ -500,3 +504,41 @@ def _is_same_type(ty: Type | TypeInfo | None, expected: TypeLike) -> bool:
             return True
 
     return False
+
+
+def get_mypy_type(node: Node) -> Type | None:
+    match node:
+        case Var(type=ty):
+            return ty
+
+        case NameExpr(node=sym):
+            match sym:
+                case Var(type=ty) | Instance(type=ty):  # type: ignore
+                    return ty
+
+                case TypeAlias(target=ty):
+                    return ty
+
+                case FuncDef(type=CallableType(ret_type=ty)):
+                    return ty
+
+        case MemberExpr(expr=lhs, name=name):
+            # TODO: don't special case this
+            match lhs:
+                case NameExpr(node=MypyFile(names=names)):
+                    match names.get(name):
+                        case SymbolTableNode(node=FuncDef(type=CallableType(ret_type=ty))):
+                            return ty
+
+            lhs_type = get_mypy_type(lhs)
+
+            if isinstance(lhs_type, Instance):
+                sym = lhs_type.type.get(name)  # type: ignore
+
+                if sym and sym.node:  # type: ignore
+                    return get_mypy_type(sym.node)  # type: ignore
+
+        case CallExpr(callee=callee):
+            return get_mypy_type(callee)
+
+    return None

refurb/checks/hashlib/use_hexdigest.py

@@ -1,8 +1,8 @@
 from dataclasses import dataclass
 
-from mypy.nodes import CallExpr, Expression, MemberExpr, NameExpr, RefExpr, Var
+from mypy.nodes import CallExpr, Expression, MemberExpr
 
-from refurb.checks.common import is_same_type, stringify
+from refurb.checks.common import get_mypy_type, is_same_type, stringify
 from refurb.error import Error
 
 
@@ -49,18 +49,13 @@ class ErrorInfo(Error):
     "hashlib.shake_128",
     "hashlib.shake_256",
     "hashlib._Hash",
+    "hashlib._BlakeHash",
+    "hashlib._VarLenHash",
 }
 
 
 def is_hashlib_algo(expr: Expression) -> bool:
-    match expr:
-        case CallExpr(callee=RefExpr(fullname=fn)) if fn in HASHLIB_ALGOS:
-            return True
-
-        case NameExpr(node=Var(type=ty)) if is_same_type(ty, *HASHLIB_ALGOS):
-            return True
-
-    return False
+    return is_same_type(get_mypy_type(expr), *HASHLIB_ALGOS)
 
 
 def check(node: CallExpr, errors: list[Error]) -> None:

test/data/err_181.py

@@ -9,6 +9,7 @@
     sha3_384,
     sha3_512,
     sha224,
+    _Hash,
 )
 from hashlib import sha256
 from hashlib import sha256 as hash_algo
@@ -40,6 +41,23 @@
 h = sha256()
 h.digest().hex()
 
+b = blake2b()
+b.digest().hex()
+
+s = shake_128()
+s.digest(10).hex()
+
+class C:
+    h: _Hash
+    b: blake2b
+
+c = C()
+c.h.digest().hex()
+c.b.digest().hex()
+
+# TODO: support this as well
+C().h.digest().hex()
+
 
 # these will not
 

test/data/err_181.txt

@@ -1,18 +1,22 @@
-test/data/err_181.py:19:1 [FURB181]: Replace `blake2b().digest().hex()` with `blake2b().hexdigest()`
-test/data/err_181.py:20:1 [FURB181]: Replace `blake2s().digest().hex()` with `blake2s().hexdigest()`
-test/data/err_181.py:21:1 [FURB181]: Replace `md5().digest().hex()` with `md5().hexdigest()`
-test/data/err_181.py:22:1 [FURB181]: Replace `sha1().digest().hex()` with `sha1().hexdigest()`
-test/data/err_181.py:23:1 [FURB181]: Replace `sha224().digest().hex()` with `sha224().hexdigest()`
-test/data/err_181.py:24:1 [FURB181]: Replace `sha256().digest().hex()` with `sha256().hexdigest()`
-test/data/err_181.py:25:1 [FURB181]: Replace `sha384().digest().hex()` with `sha384().hexdigest()`
-test/data/err_181.py:26:1 [FURB181]: Replace `sha3_224().digest().hex()` with `sha3_224().hexdigest()`
-test/data/err_181.py:27:1 [FURB181]: Replace `sha3_256().digest().hex()` with `sha3_256().hexdigest()`
-test/data/err_181.py:28:1 [FURB181]: Replace `sha3_384().digest().hex()` with `sha3_384().hexdigest()`
-test/data/err_181.py:29:1 [FURB181]: Replace `sha3_512().digest().hex()` with `sha3_512().hexdigest()`
-test/data/err_181.py:30:1 [FURB181]: Replace `sha512().digest().hex()` with `sha512().hexdigest()`
-test/data/err_181.py:31:1 [FURB181]: Replace `shake_128().digest(10).hex()` with `shake_128().hexdigest(10)`
-test/data/err_181.py:32:1 [FURB181]: Replace `shake_256().digest(10).hex()` with `shake_256().hexdigest(10)`
-test/data/err_181.py:34:1 [FURB181]: Replace `hashlib.sha256().digest().hex()` with `hashlib.sha256().hexdigest()`
-test/data/err_181.py:36:1 [FURB181]: Replace `sha256(b"text").digest().hex()` with `sha256(b"text").hexdigest()`
-test/data/err_181.py:38:1 [FURB181]: Replace `hash_algo().digest().hex()` with `hash_algo().hexdigest()`
-test/data/err_181.py:41:1 [FURB181]: Replace `h.digest().hex()` with `h.hexdigest()`
+test/data/err_181.py:20:1 [FURB181]: Replace `blake2b().digest().hex()` with `blake2b().hexdigest()`
+test/data/err_181.py:21:1 [FURB181]: Replace `blake2s().digest().hex()` with `blake2s().hexdigest()`
+test/data/err_181.py:22:1 [FURB181]: Replace `md5().digest().hex()` with `md5().hexdigest()`
+test/data/err_181.py:23:1 [FURB181]: Replace `sha1().digest().hex()` with `sha1().hexdigest()`
+test/data/err_181.py:24:1 [FURB181]: Replace `sha224().digest().hex()` with `sha224().hexdigest()`
+test/data/err_181.py:25:1 [FURB181]: Replace `sha256().digest().hex()` with `sha256().hexdigest()`
+test/data/err_181.py:26:1 [FURB181]: Replace `sha384().digest().hex()` with `sha384().hexdigest()`
+test/data/err_181.py:27:1 [FURB181]: Replace `sha3_224().digest().hex()` with `sha3_224().hexdigest()`
+test/data/err_181.py:28:1 [FURB181]: Replace `sha3_256().digest().hex()` with `sha3_256().hexdigest()`
+test/data/err_181.py:29:1 [FURB181]: Replace `sha3_384().digest().hex()` with `sha3_384().hexdigest()`
+test/data/err_181.py:30:1 [FURB181]: Replace `sha3_512().digest().hex()` with `sha3_512().hexdigest()`
+test/data/err_181.py:31:1 [FURB181]: Replace `sha512().digest().hex()` with `sha512().hexdigest()`
+test/data/err_181.py:32:1 [FURB181]: Replace `shake_128().digest(10).hex()` with `shake_128().hexdigest(10)`
+test/data/err_181.py:33:1 [FURB181]: Replace `shake_256().digest(10).hex()` with `shake_256().hexdigest(10)`
+test/data/err_181.py:35:1 [FURB181]: Replace `hashlib.sha256().digest().hex()` with `hashlib.sha256().hexdigest()`
+test/data/err_181.py:37:1 [FURB181]: Replace `sha256(b"text").digest().hex()` with `sha256(b"text").hexdigest()`
+test/data/err_181.py:39:1 [FURB181]: Replace `hash_algo().digest().hex()` with `hash_algo().hexdigest()`
+test/data/err_181.py:42:1 [FURB181]: Replace `h.digest().hex()` with `h.hexdigest()`
+test/data/err_181.py:45:1 [FURB181]: Replace `b.digest().hex()` with `b.hexdigest()`
+test/data/err_181.py:48:1 [FURB181]: Replace `s.digest(10).hex()` with `s.hexdigest(10)`
+test/data/err_181.py:55:1 [FURB181]: Replace `c.h.digest().hex()` with `c.h.hexdigest()`
+test/data/err_181.py:56:1 [FURB181]: Replace `c.b.digest().hex()` with `c.b.hexdigest()`