diff --git a/src/Forms/GridFieldConfig_Lumberjack.php b/src/Forms/GridFieldConfig_Lumberjack.php
index 1330a40..443097f 100644
--- a/src/Forms/GridFieldConfig_Lumberjack.php
+++ b/src/Forms/GridFieldConfig_Lumberjack.php
@@ -10,9 +10,6 @@
 use SilverStripe\Forms\GridField\GridFieldPaginator;
 use SilverStripe\Forms\GridField\GridFieldSortableHeader;
 use SilverStripe\Forms\GridField\GridFieldToolbarHeader;
-use SilverStripe\Lumberjack\Forms\GridFieldSiteTreeAddNewButton;
-use SilverStripe\Lumberjack\Forms\GridFieldSiteTreeEditButton;
-use SilverStripe\Lumberjack\Forms\GridFieldSiteTreeState;
 
 /**
  * GridField config necessary for managing a SiteTree object.
diff --git a/src/Forms/GridFieldSiteTreeEditButton.php b/src/Forms/GridFieldSiteTreeEditButton.php
index a3f6de6..ed5825f 100644
--- a/src/Forms/GridFieldSiteTreeEditButton.php
+++ b/src/Forms/GridFieldSiteTreeEditButton.php
@@ -2,7 +2,9 @@
 
 namespace SilverStripe\Lumberjack\Forms;
 
+use SilverStripe\Forms\GridField\GridField;
 use SilverStripe\Forms\GridField\GridFieldEditButton;
+use SilverStripe\ORM\DataObject;
 use SilverStripe\View\ArrayData;
 
 /**
@@ -10,11 +12,8 @@
  *
  * Bypasses GridFieldDetailForm
  *
- * @package silverstripe
- * @subpackage lumberjack
- *
  * @author Michael Strong <mstrong@silverstripe.org>
-**/
+ **/
 class GridFieldSiteTreeEditButton extends GridFieldEditButton
 {
     /**
@@ -29,8 +28,8 @@ public function getColumnContent($gridField, $record, $columnName)
         // which can make the form readonly if no edit permissions are available.
 
         $data = ArrayData::create([
-              'Link' => $record->CMSEditLink(),
-              'ExtraClass' => $this->getExtraClass()
+            'Link' => $record->CMSEditLink(),
+            'ExtraClass' => $this->getExtraClass(),
         ]);
 
         return $data->renderWith(GridFieldEditButton::class);