CLN: only construct case necessary for benchmark

pandas-dev · Jan 24, 2019 · d8bfcef · d8bfcef
1 parent 9c8d273
commit d8bfcef
Showing 1 changed file with 10 additions and 9 deletions.
diff --git a/asv_bench/benchmarks/categoricals.py b/asv_bench/benchmarks/categoricals.py
@@ -225,15 +225,16 @@ def setup(self, index):
         N = 10**6
         categories = ['a', 'b', 'c']
         values = [0] * N + [1] * N + [2] * N
-        indices = {
-            'monotonic_incr': pd.Categorical.from_codes(values,
-                                                        categories=categories),
-            'monotonic_decr': pd.Categorical.from_codes(list(reversed(values)),
-                                                        categories=categories),
-            'non_monotonic': pd.Categorical.from_codes([0, 1, 2] * N,
-                                                       categories=categories)
-        }
-        self.data = indices[index]
+        if index == 'monotonic_incr':
+            self.data = pd.Categorical.from_codes(values, categories=categories)
+        elif index == 'monotonic_decr':
+            self.data = pd.Categorical.from_codes(list(reversed(values)),
+                                                        categories=categories)
+        elif index == 'non_monotonic':
+            self.data = pd.Categorical.from_codes([0, 1, 2] * N,
+                                                  categories=categories)
+        else:
+            raise ValueError('Invalid index param: {}'.format(index))
 
         self.scalar = 10000
         self.list = list(range(10000))