diff --git a/modules/features_ui/src/Form/FeaturesDiffForm.php b/modules/features_ui/src/Form/FeaturesDiffForm.php
index 34febf9..a55703b 100644
--- a/modules/features_ui/src/Form/FeaturesDiffForm.php
+++ b/modules/features_ui/src/Form/FeaturesDiffForm.php
@@ -13,6 +13,7 @@ use Drupal\features\FeaturesAssignerInterface;
 use Drupal\features\FeaturesManagerInterface;
 use Drupal\Core\Form\FormBase;
 use Drupal\Core\Form\FormStateInterface;
+use Drupal\features\Package;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 use Drupal\Component\Diff\DiffFormatter;
 use Drupal\config_update\ConfigRevertInterface;
@@ -203,8 +204,8 @@ class FeaturesDiffForm extends FormBase {
   /**
    * Returns a form element for the given overrides.
    *
-   * @param array $package
-   *   A package array.
+   * @param \Drupal\features\Package $package
+   *   A package.
    * @param array $overrides
    *   An array of overrides.
    * @param array $missing
@@ -213,7 +214,7 @@ class FeaturesDiffForm extends FormBase {
    * @return array
    *   A form element.
    */
-  protected function diffOutput($package, $overrides, $missing = array()) {
+  protected function diffOutput(Package $package, $overrides, $missing = array()) {
     $element = array();
     $config = $this->featuresManager->getConfigCollection();
     $components = array_merge($missing, $overrides);
@@ -263,7 +264,7 @@ class FeaturesDiffForm extends FormBase {
           ),
         ),
         '#attributes' => array(
-          'class' => 'diff-' . $package['machine_name'],
+          'class' => 'diff-' . $package->getMachineName(),
         ),
       );
     }
diff --git a/modules/features_ui/src/Form/FeaturesEditForm.php b/modules/features_ui/src/Form/FeaturesEditForm.php
index a01e25b..d0713a6 100644
--- a/modules/features_ui/src/Form/FeaturesEditForm.php
+++ b/modules/features_ui/src/Form/FeaturesEditForm.php
@@ -602,6 +602,7 @@ class FeaturesEditForm extends FormBase {
     }
 
     // Generate new populated feature.
+    $export['package'] = $this->package;
     $export['config_new'] = $config_new;
 
     // Now fill the $export with categorized sections of component options
diff --git a/modules/features_ui/src/Form/FeaturesExportForm.php b/modules/features_ui/src/Form/FeaturesExportForm.php
index 70c1370..fc33756 100644
--- a/modules/features_ui/src/Form/FeaturesExportForm.php
+++ b/modules/features_ui/src/Form/FeaturesExportForm.php
@@ -284,7 +284,7 @@ class FeaturesExportForm extends FormBase {
    * Builds the details of a package.
    *
    * @param \Drupal\features\Package $package
-   *   The package name.
+   *   The package.
    *
    * @return array
    *   A render array of a form element.
diff --git a/src/FeaturesGenerator.php b/src/FeaturesGenerator.php
index 3d94861..fb26287 100644
--- a/src/FeaturesGenerator.php
+++ b/src/FeaturesGenerator.php
@@ -148,13 +148,13 @@ class FeaturesGenerator implements FeaturesGeneratorInterface {
         // The install profile doesn't need renaming.
         if ($package['type'] != 'profile') {
           unset($packages[$package_name]);
-          $package['machine_name'] = $bundle->getFullName($package['machine_name']);
-          $packages[$package['machine_name']] = $package;
+          $package->setMachineName($bundle->getFullName($package->getMachineName()));
+          $packages[$package->getMachineName()] = $package;
         }
 
         // Set the bundle machine name.
-        $packages[$package['machine_name']]['bundle'] = $bundle->getMachineName();
-        $new_package_names[] = $package['machine_name'];
+        $packages[$package->getMachineName()]['bundle'] = $bundle->getMachineName();
+        $new_package_names[] = $package->getMachineName();
       }
       $this->featuresManager->setPackages($packages);
       $package_names = $new_package_names;