diff --git a/src/Extension/MemberExtension.php b/src/Extension/MemberExtension.php
index 97afef4..5cf29a6 100644
--- a/src/Extension/MemberExtension.php
+++ b/src/Extension/MemberExtension.php
@@ -3,9 +3,10 @@
 namespace SilverStripe\RealMe\Extension;
 
 use SilverStripe\ORM\DataExtension;
+use SilverStripe\Security\Member;
 
 /**
- * Class MemberExtension
+ * @extends DataExtension<Member>
  */
 class MemberExtension extends DataExtension
 {
diff --git a/src/Extension/SiteTreeExtension.php b/src/Extension/SiteTreeExtension.php
index 8cbeedc..42b451a 100644
--- a/src/Extension/SiteTreeExtension.php
+++ b/src/Extension/SiteTreeExtension.php
@@ -2,12 +2,16 @@
 
 namespace SilverStripe\RealMe\Extension;
 
+use SilverStripe\CMS\Model\SiteTree;
 use SilverStripe\RealMe\RealMeService;
 use SilverStripe\Security\InheritedPermissions;
 use SilverStripe\Security\Member;
 use SilverStripe\ORM\DataObject;
 use SilverStripe\ORM\DataExtension;
 
+/**
+ * @extends DataExtension<SiteTree>
+ */
 class SiteTreeExtension extends DataExtension
 {
     private static $dependencies = array(
@@ -37,7 +41,7 @@ public function canView($member)
         if ($member && $member->ID) {
             return null;
         }
-        
+
         $data = $this->service->getUserData();
         if (empty($data)) {
             // Defer if there's no logged in RealMe user