Add NatSpec on state variables

Author: smonicas

crytic_compile/utils/natspec.py

@@ -2,6 +2,64 @@
 Natspec module https://solidity.readthedocs.io/en/latest/natspec-format.html
 """
 
+class DevStateVariable:
+    """
+    Model the dev state variable
+    """
+
+    def __init__(self, variable: dict) -> None:
+        """Init the object
+
+        Args:
+            method (Dict): Method infos (details, params, returns, custom:*)
+        """
+        self._details: str | None = variable.get("details", None)
+        # For state variable there can be only 1 return documentation.
+        # We only keep the value that is the documentation, the key is always _0 
+        self._return: str | None = variable["returns"]["_0"] if "returns" in variable else None
+        # Extract custom fields (keys starting with "custom:")
+        self._custom: dict[str, str] = {k: v for k, v in variable.items() if k.startswith("custom:")}
+
+    @property
+    def details(self) -> str | None:
+        """Return the state variable details
+
+        Returns:
+            Optional[str]: state variable details
+        """
+        return self._details
+
+    @property
+    def variable_return(self) -> str | None:
+        """Return the state variable returns
+
+        Returns:
+            Optional[str]: state variable returns
+        """
+        return self._return
+
+    @property
+    def custom(self) -> dict[str, str]:
+        """Return the state variable custom fields
+
+        Returns:
+            Dict[str, str]: custom field name => value (e.g. "custom:security" => "value")
+        """
+        return self._custom
+
+    def export(self) -> dict:
+        """Export to a python dict
+
+        Returns:
+            Dict: Exported dev state variable
+        """
+        result = {
+            "details": self.details,
+            "return": self.variable_return,
+            "custom": self.custom,
+        }
+        return result
+
 
 class UserMethod:
     """
@@ -179,6 +237,9 @@ def __init__(self, devdoc: dict):
         self._methods: dict[str, DevMethod] = {
             k: DevMethod(item) for k, item in devdoc.get("methods", {}).items()
         }
+        self._state_variables: dict[str, DevStateVariable] = {
+            k: DevStateVariable(item) for k, item in devdoc.get("stateVariables", {}).items()
+        }
         self._title: str | None = devdoc.get("title", None)
         # Extract contract-level custom fields (keys starting with "custom:")
         self._custom: dict[str, str] = {k: v for k, v in devdoc.items() if k.startswith("custom:")}
@@ -210,6 +271,15 @@ def methods(self) -> dict[str, DevMethod]:
         """
         return self._methods
 
+    @property
+    def state_variables(self) -> dict[str, DevStateVariable]:
+        """Return the dev state variables
+
+        Returns:
+            Dict[str, DevStateVariable]: state_variable_name => DevStateVariable
+        """
+        return self._state_variables
+
     @property
     def title(self) -> str | None:
         """Return the dev title
@@ -240,6 +310,7 @@ def export(self) -> dict:
             "details": self.details,
             "title": self.title,
             "custom": self.custom,
+            "state_variables": self.state_variables
         }
         return result
 

tests/test_natspec.py

@@ -2,7 +2,14 @@
 Test NatSpec parsing, including custom fields (@custom:*)
 """
 
-from crytic_compile.utils.natspec import DevDoc, DevMethod, Natspec, UserDoc, UserMethod
+from crytic_compile.utils.natspec import (
+    DevDoc,
+    DevMethod,
+    DevStateVariable,
+    Natspec,
+    UserDoc,
+    UserMethod,
+)
 
 
 class TestUserMethod:
@@ -94,6 +101,73 @@ def test_devmethod_empty_method(self) -> None:
         assert method.custom == {}
 
 
+class TestDevStateVariable:
+    """Tests for DevStateVariable class"""
+
+    def test_devstatevariable_basic_fields(self) -> None:
+        """Test DevStateVariable parses basic fields"""
+        var_data = {
+            "details": "The total supply of tokens",
+            "returns": {"_0": "The total token supply"},
+        }
+        var = DevStateVariable(var_data)
+        assert var.details == "The total supply of tokens"
+        assert var.variable_return == "The total token supply"
+
+    def test_devstatevariable_custom_fields(self) -> None:
+        """Test DevStateVariable extracts custom fields"""
+        var_data = {
+            "details": "Owner address",
+            "returns": {"_0": "The contract owner"},
+            "custom:security": "critical",
+            "custom:access": "public",
+        }
+        var = DevStateVariable(var_data)
+        assert var.custom == {
+            "custom:security": "critical",
+            "custom:access": "public",
+        }
+
+    def test_devstatevariable_no_custom_fields(self) -> None:
+        """Test DevStateVariable returns empty dict when no custom fields"""
+        var_data = {
+            "details": "Balance mapping",
+            "returns": {"_0": "User balance"},
+        }
+        var = DevStateVariable(var_data)
+        assert var.custom == {}
+
+    def test_devstatevariable_no_returns(self) -> None:
+        """Test DevStateVariable with no returns field"""
+        var_data = {
+            "details": "Some variable",
+        }
+        var = DevStateVariable(var_data)
+        assert var.details == "Some variable"
+        assert var.variable_return is None
+
+    def test_devstatevariable_export(self) -> None:
+        """Test DevStateVariable export"""
+        var_data = {
+            "details": "Token name",
+            "returns": {"_0": "The token name string"},
+            "custom:immutable": "true",
+        }
+        var = DevStateVariable(var_data)
+        exported = var.export()
+
+        assert exported["details"] == "Token name"
+        assert exported["return"] == "The token name string"
+        assert exported["custom"]["custom:immutable"] == "true"
+
+    def test_devstatevariable_empty(self) -> None:
+        """Test DevStateVariable with empty dict"""
+        var = DevStateVariable({})
+        assert var.details is None
+        assert var.variable_return is None
+        assert var.custom == {}
+
+
 class TestUserDoc:
     """Tests for UserDoc class"""
 
@@ -230,6 +304,47 @@ def test_devdoc_export_with_methods_custom(self) -> None:
         assert exported["methods"]["foo()"]["custom"]["custom:audit"] == "verified"
         assert exported["methods"]["foo()"]["details"] == "foo details"
 
+    def test_devdoc_state_variables(self) -> None:
+        """Test DevDoc parses stateVariables"""
+        devdoc_data = {
+            "title": "Test Contract",
+            "methods": {},
+            "stateVariables": {
+                "totalSupply": {
+                    "details": "Total token supply",
+                    "returns": {"_0": "The total supply"},
+                },
+                "owner": {
+                    "details": "Contract owner",
+                    "returns": {"_0": "Owner address"},
+                    "custom:security": "critical",
+                },
+            },
+        }
+        devdoc = DevDoc(devdoc_data)
+
+        assert "totalSupply" in devdoc.state_variables
+        assert "owner" in devdoc.state_variables
+
+        total_supply = devdoc.state_variables["totalSupply"]
+        assert total_supply.details == "Total token supply"
+        assert total_supply.variable_return == "The total supply"
+        assert total_supply.custom == {}
+
+        owner = devdoc.state_variables["owner"]
+        assert owner.details == "Contract owner"
+        assert owner.variable_return == "Owner address"
+        assert owner.custom == {"custom:security": "critical"}
+
+    def test_devdoc_no_state_variables(self) -> None:
+        """Test DevDoc with no stateVariables"""
+        devdoc_data = {
+            "title": "Test Contract",
+            "methods": {},
+        }
+        devdoc = DevDoc(devdoc_data)
+        assert devdoc.state_variables == {}
+
 
 class TestNatspec:
     """Tests for Natspec class"""