Drop variables in a batch in normalize_entity

featuretools/entityset/entity.py

@@ -401,14 +401,16 @@ def add_interesting_values(self, max_values=5, verbose=False):
 
         self.entityset.reset_data_description()
 
-    def delete_variable(self, variable_id):
+    def delete_variables(self, variable_ids):
         """
-        Remove variable from entity's dataframe and from
+        Remove variables from entity's dataframe and from
         self.variables
         """
-        self.df.drop(variable_id, axis=1, inplace=True)
-        v = self._get_variable(variable_id)
-        self.variables.remove(v)
+        self.df.drop(variable_ids, axis=1, inplace=True)
+
+        for v_id in variable_ids:
+            v = self._get_variable(v_id)
+            self.variables.remove(v)
 
     def set_time_index(self, variable_id, already_sorted=False):
         # check time type

featuretools/entityset/entityset.py

@@ -794,8 +794,7 @@ def normalize_entity(self, base_entity_id, new_entity_id, index,
             secondary_time_index=make_secondary_time_index,
             variable_types=transfer_types)
 
-        for v in additional_variables:
-            self.entity_dict[base_entity_id].delete_variable(v)
+        self.entity_dict[base_entity_id].delete_variables(additional_variables)
 
         new_entity = self.entity_dict[new_entity_id]
         base_entity.convert_variable_type(base_entity_index, vtypes.Id, convert_data=False)

featuretools/tests/entityset_tests/test_entity.py

@@ -112,3 +112,15 @@ def test_query_by_values_returns_rows_in_given_order():
                                   })
     query = es['test'].query_by_values(['b', 'a'], variable_id='value')
     assert np.array_equal(query['id'], [1, 3, 4, 5])
+
+
+def test_delete_variables(es):
+    entity = es['customers']
+    to_delete = ['age', 'cohort', 'email']
+    entity.delete_variables(to_delete)
+
+    variable_names = [v.id for v in entity.variables]
+
+    for var in to_delete:
+        assert var not in variable_names
+        assert var not in entity.df