diff --git a/lib/internal/Magento/Framework/Code/Generator/ClassGenerator.php b/lib/internal/Magento/Framework/Code/Generator/ClassGenerator.php
index 6978242842a3f..677f81ecca312 100644
--- a/lib/internal/Magento/Framework/Code/Generator/ClassGenerator.php
+++ b/lib/internal/Magento/Framework/Code/Generator/ClassGenerator.php
@@ -61,27 +61,6 @@ class ClassGenerator extends \Zend\Code\Generator\ClassGenerator implements
         'passedByReference' => 'setPassedByReference',
     ];
 
-    /**
-     * @var array Array of string names
-     */
-    protected $traits = array();
-
-    public function setTraits(array $traits)
-    {
-        $this->traits = $traits;
-        return $this;
-    }
-
-    /**
-     * Returns the "traits" classes
-     *
-     * @return array
-     */
-    public function getTraits()
-    {
-        return $this->traits;
-    }
-
     /**
      * @param object $object
      * @param array $data
@@ -230,77 +209,4 @@ public function getNamespaceName()
     {
         return ltrim(parent::getNamespaceName(), '\\');
     }
-
-    /**
-     * @return string
-     */
-    public function generate()
-    {
-        if (!$this->isSourceDirty()) {
-            $output = $this->getSourceContent();
-            if (!empty($output)) {
-                return $output;
-            }
-        }
-
-        $output = '';
-
-        if (null !== ($namespace = $this->getNamespaceName())) {
-            $output .= 'namespace ' . $namespace . ';' . self::LINE_FEED . self::LINE_FEED;
-        }
-
-        $uses = $this->getUses();
-        if (!empty($uses)) {
-            foreach ($uses as $use) {
-                $output .= 'use ' . $use . ';' . self::LINE_FEED;
-            }
-            $output .= self::LINE_FEED;
-        }
-
-        if (null !== ($docBlock = $this->getDocBlock())) {
-            $docBlock->setIndentation('');
-            $output .= $docBlock->generate();
-        }
-
-        if ($this->isAbstract()) {
-            $output .= 'abstract ';
-        }
-
-        $output .= 'class ' . $this->getName();
-
-        if (!empty($this->extendedClass)) {
-            $output .= ' extends ' . $this->extendedClass;
-        }
-
-        $implemented = $this->getImplementedInterfaces();
-        if (!empty($implemented)) {
-            $output .= ' implements ' . implode(', ', $implemented);
-        }
-
-        $output .= self::LINE_FEED . '{' . self::LINE_FEED . self::LINE_FEED;
-
-        $traits = $this->getTraits();
-        if (!empty($traits)) {
-            $output .= self::LINE_FEED . $this->indentation
-                . 'use ' . (implode(', ', $traits)) . ';' . self::LINE_FEED . self::LINE_FEED;
-        }
-
-        $properties = $this->getProperties();
-        if (!empty($properties)) {
-            foreach ($properties as $property) {
-                $output .= $property->generate() . self::LINE_FEED . self::LINE_FEED;
-            }
-        }
-
-        $methods = $this->getMethods();
-        if (!empty($methods)) {
-            foreach ($methods as $method) {
-                $output .= $method->generate() . self::LINE_FEED;
-            }
-        }
-
-        $output .= self::LINE_FEED . '}' . self::LINE_FEED;
-
-        return $output;
-    }
 }
diff --git a/lib/internal/Magento/Framework/Code/Generator/CodeGeneratorInterface.php b/lib/internal/Magento/Framework/Code/Generator/CodeGeneratorInterface.php
index 617f1775582fe..c708bad2ca6bf 100644
--- a/lib/internal/Magento/Framework/Code/Generator/CodeGeneratorInterface.php
+++ b/lib/internal/Magento/Framework/Code/Generator/CodeGeneratorInterface.php
@@ -56,10 +56,10 @@ public function setExtendedClass($extendedClass);
     public function setImplementedInterfaces(array $interfaces);
 
     /**
-     * Set a list of traits.
+     * Add a trait to the class.
      *
-     * @param array $traits
+     * @param $trait
      * @return $this
      */
-    public function setTraits(array $traits);
+    public function addTrait($trait);
 }
diff --git a/lib/internal/Magento/Framework/Interception/Code/Generator/Interceptor.php b/lib/internal/Magento/Framework/Interception/Code/Generator/Interceptor.php
index eda4a5c20ccd7..413188abd75f1 100644
--- a/lib/internal/Magento/Framework/Interception/Code/Generator/Interceptor.php
+++ b/lib/internal/Magento/Framework/Interception/Code/Generator/Interceptor.php
@@ -161,7 +161,7 @@ protected function _generateCode()
         } else {
             $this->_classGenerator->setExtendedClass($typeName);
         }
-        $this->_classGenerator->setTraits(['\Magento\Framework\Interception\Interceptor']);
+        $this->_classGenerator->addTrait('\Magento\Framework\Interception\Interceptor');
         return parent::_generateCode();
     }