Fix check_contracts handling of forward references in type annotations (

#941)

CHANGELOG.md

@@ -11,6 +11,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 - Make `graphviz` an optional dependency, and clarify the installation requirements for visualizing
  control flow graphs.
+- Fix `check_contrats` handling of forward references in class type annotations when using `check_contracts` decorator.
 
 ## [2.6.0] - 2023-08-06
 

python_ta/contracts/__init__.py

@@ -140,7 +140,7 @@ def add_class_invariants(klass: type) -> None:
  _set_invariants(klass)
 
  klass_mod = _get_module(klass)
- cls_annotations = typing.get_type_hints(klass, localns=klass_mod.__dict__)
+ cls_annotations = None # This is a cached value set the first time new_setattr is called
 
  def new_setattr(self: klass, name: str, value: Any) -> None:
  """Set the value of the given attribute on self to the given value.
@@ -150,6 +150,11 @@ def new_setattr(self: klass, name: str, value: Any) -> None:
  if not ENABLE_CONTRACT_CHECKING:
  super(klass, self).__setattr__(name, value)
  return
+
+ nonlocal cls_annotations
+ if cls_annotations is None:
+ cls_annotations = typing.get_type_hints(klass, localns=klass_mod.__dict__)
+
  if name in cls_annotations:
  try:
  _debug(f"Checking type of attribute {attr} for {klass.__qualname__} instance")

tests/test_contracts/test_class_forward_reference.py

@@ -0,0 +1,53 @@
+"""Tests for forward references in instance attribute type annotations."""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import pytest
+
+from python_ta.contracts import check_contracts
+
+
+@check_contracts
+class Node:
+ """A node in a linked list."""
+
+ item: Any
+ next: Optional[Node]
+
+ def __init__(self, item: Any) -> None:
+ self.item = item
+ self.next = None
+
+ def bad_method(self) -> None:
+ """Set self.next to an invalid value, violating the type annotation."""
+ self.next = 1
+
+
+def test_node_no_error() -> None:
+ """Test that a Node can be initialized without error."""
+ Node(1)
+
+
+def test_node_setattr_no_error() -> None:
+ """Test that assigning a valid attribute to a Node does not raise an error."""
+ node = Node(1)
+ node.next = Node(2)
+
+ assert node.next.item == 2
+
+
+def test_node_setattr_error() -> None:
+ """Test that assigning an invalid attribute to a Node raises an AssertionError."""
+ node = Node(1)
+
+ with pytest.raises(AssertionError):
+ node.next = 2
+
+
+def test_node_method_error() -> None:
+ """Test that violating a Node type annotation with a method call raises an AssertionError."""
+ node = Node(1)
+
+ with pytest.raises(AssertionError):
+ node.bad_method()

tests/test_contracts_type_alias_abstract_ring.py → tests/test_contracts/test_contracts_type_alias_abstract_ring.py

@@ -1,5 +1,7 @@
 from python_ta.contracts import check_contracts
-from tests.test_contracts_type_alias_abstract_network import AbstractNetwork
+from tests.test_contracts.test_contracts_type_alias_abstract_network import (
+ AbstractNetwork,
+)
 
 
 def test_type_alias_as_type_annotation_for_class_attribute_no_error() -> None: