Add contracts to DjangoXBlockUserStateClient interface methods

lms/djangoapps/courseware/user_state_client.py

@@ -11,9 +11,13 @@ try:
 except ImportError:
     import json
 
-from xblock.fields import Scope
+from xblock.fields import Scope, ScopeBase
 from xblock_user_state.interface import XBlockUserStateClient
 from courseware.models import StudentModule
+from contracts import contract, new_contract
+from opaque_keys.edx.keys import UsageKey
+
+new_contract('UsageKey', UsageKey)
 
 
 class DjangoXBlockUserStateClient(XBlockUserStateClient):
@@ -42,6 +46,12 @@ class DjangoXBlockUserStateClient(XBlockUserStateClient):
     def __init__(self, user):
         self.user = user
 
+    @contract(
+        username="basestring",
+        block_key=UsageKey,
+        scope=ScopeBase,
+        fields="seq(basestring)|set(basestring)|None"
+    )
     def get(self, username, block_key, scope=Scope.user_state, fields=None):
         """
         Retrieve the stored XBlock state for a single xblock usage.
@@ -66,6 +76,7 @@ class DjangoXBlockUserStateClient(XBlockUserStateClient):
 
         return state
 
+    @contract(username="basestring", block_key=UsageKey, state="dict(basestring: *)", scope=ScopeBase)
     def set(self, username, block_key, state, scope=Scope.user_state):
         """
         Set fields for a particular XBlock.
@@ -79,6 +90,12 @@ class DjangoXBlockUserStateClient(XBlockUserStateClient):
         assert self.user.username == username
         self.set_many(username, {block_key: state}, scope)
 
+    @contract(
+        username="basestring",
+        block_key=UsageKey,
+        scope=ScopeBase,
+        fields="seq(basestring)|set(basestring)|None"
+    )
     def delete(self, username, block_key, scope=Scope.user_state, fields=None):
         """
         Delete the stored XBlock state for a single xblock usage.
@@ -92,7 +109,12 @@ class DjangoXBlockUserStateClient(XBlockUserStateClient):
         assert self.user.username == username
         return self.delete_many(username, [block_key], scope, fields=fields)
 
-    @contract(username="basestring", block_key=UsageKey, scope=ScopeBase, fields="seq(basestring)|set(basestring)|None")
+    @contract(
+        username="basestring",
+        block_key=UsageKey,
+        scope=ScopeBase,
+        fields="seq(basestring)|set(basestring)|None"
+    )
     def get_mod_date(self, username, block_key, scope=Scope.user_state, fields=None):
         """
         Get the last modification date for fields from the specified blocks.
@@ -140,7 +162,12 @@ class DjangoXBlockUserStateClient(XBlockUserStateClient):
                 usage_key = student_module.module_state_key.map_into_course(student_module.course_id)
                 yield (student_module, usage_key)
 
-
+    @contract(
+        username="basestring",
+        block_keys="seq(UsageKey)|set(UsageKey)",
+        scope=ScopeBase,
+        fields="seq(basestring)|set(basestring)|None"
+    )
     def get_many(self, username, block_keys, scope=Scope.user_state, fields=None):
         """
         Retrieve the stored XBlock state for a single xblock usage.
@@ -167,6 +194,7 @@ class DjangoXBlockUserStateClient(XBlockUserStateClient):
                 state = json.loads(module.state)
             yield (usage_key, state)
 
+    @contract(username="basestring", block_keys_to_state="dict(UsageKey: dict(basestring: *))", scope=ScopeBase)
     def set_many(self, username, block_keys_to_state, scope=Scope.user_state):
         """
         Set fields for a particular XBlock.
@@ -208,6 +236,12 @@ class DjangoXBlockUserStateClient(XBlockUserStateClient):
                 # We just read this object, so we know that we can do an update
                 student_module.save(force_update=True)
 
+    @contract(
+        username="basestring",
+        block_keys="seq(UsageKey)|set(UsageKey)",
+        scope=ScopeBase,
+        fields="seq(basestring)|set(basestring)|None"
+    )
     def delete_many(self, username, block_keys, scope=Scope.user_state, fields=None):
         """
         Delete the stored XBlock state for a many xblock usages.