From d73048ef7d284aa240254324154e7fe18d1571a0 Mon Sep 17 00:00:00 2001
From: Nikita Titov <nekit94-08@mail.ru>
Date: Mon, 4 May 2020 17:56:51 +0300
Subject: [PATCH] Use list comprehensions where possible (#201)

---
 m2cgen/assemblers/boosting.py         |  6 +-----
 m2cgen/assemblers/ensemble.py         |  3 +--
 m2cgen/assemblers/linear.py           | 17 ++++++++---------
 m2cgen/assemblers/svm.py              | 15 ++++++++-------
 m2cgen/assemblers/utils.py            |  7 ++++---
 m2cgen/interpreters/code_generator.py |  9 ++++-----
 6 files changed, 26 insertions(+), 31 deletions(-)

diff --git a/m2cgen/assemblers/boosting.py b/m2cgen/assemblers/boosting.py
index 714da1b6..93613c4a 100644
--- a/m2cgen/assemblers/boosting.py
+++ b/m2cgen/assemblers/boosting.py
@@ -245,8 +245,4 @@ def _assemble_tree(self, tree):
 def _split_estimator_params_by_classes(values, n_classes):
     # Splits are computed based on a comment
     # https://github.com/dmlc/xgboost/issues/1746#issuecomment-267400592.
-    estimator_params_by_classes = [[] for _ in range(n_classes)]
-    for i in range(len(values)):
-        class_idx = i % n_classes
-        estimator_params_by_classes[class_idx].append(values[i])
-    return estimator_params_by_classes
+    return [values[class_idx::n_classes] for class_idx in range(n_classes)]
diff --git a/m2cgen/assemblers/ensemble.py b/m2cgen/assemblers/ensemble.py
index 7437a02c..694ed9ba 100644
--- a/m2cgen/assemblers/ensemble.py
+++ b/m2cgen/assemblers/ensemble.py
@@ -1,7 +1,6 @@
 from m2cgen import ast
-from m2cgen.assemblers import utils
+from m2cgen.assemblers import utils, TreeModelAssembler
 from m2cgen.assemblers.base import ModelAssembler
-from m2cgen.assemblers import TreeModelAssembler
 
 
 class RandomForestModelAssembler(ModelAssembler):
diff --git a/m2cgen/assemblers/linear.py b/m2cgen/assemblers/linear.py
index c50d58c7..edc31dae 100644
--- a/m2cgen/assemblers/linear.py
+++ b/m2cgen/assemblers/linear.py
@@ -17,9 +17,10 @@ def _build_ast(self):
         if coef.shape[0] == 1:
             return _linear_to_ast(coef[0], intercept[0])
 
-        exprs = []
-        for idx in range(coef.shape[0]):
-            exprs.append(_linear_to_ast(coef[idx], intercept[idx]))
+        exprs = [
+            _linear_to_ast(coef[idx], intercept[idx])
+            for idx in range(coef.shape[0])
+        ]
         return ast.VectorVal(exprs)
 
     def _get_intercept(self):
@@ -71,12 +72,10 @@ def _get_coef(self):
 
 
 def _linear_to_ast(coef, intercept):
-    feature_weight_mul_ops = []
-
-    for index, value in enumerate(coef):
-        feature_weight_mul_ops.append(
-            utils.mul(ast.FeatureRef(index), ast.NumVal(value)))
-
+    feature_weight_mul_ops = [
+        utils.mul(ast.FeatureRef(index), ast.NumVal(value))
+        for index, value in enumerate(coef)
+    ]
     return utils.apply_op_to_expressions(
         ast.BinNumOpType.ADD,
         ast.NumVal(intercept),
diff --git a/m2cgen/assemblers/svm.py b/m2cgen/assemblers/svm.py
index 18ade78e..22e2e331 100644
--- a/m2cgen/assemblers/svm.py
+++ b/m2cgen/assemblers/svm.py
@@ -36,10 +36,10 @@ def _assemble_single_output(self, idx=0):
 
         kernel_exprs = self._apply_kernel(support_vectors)
 
-        kernel_weight_mul_ops = []
-        for index, value in enumerate(coef):
-            kernel_weight_mul_ops.append(
-                utils.mul(kernel_exprs[index], ast.NumVal(value)))
+        kernel_weight_mul_ops = [
+            utils.mul(kernel_exprs[index], ast.NumVal(value))
+            for index, value in enumerate(coef)
+        ]
 
         return utils.apply_op_to_expressions(
             ast.BinNumOpType.ADD,
@@ -188,9 +188,10 @@ def _get_output_size(self):
         return output_size
 
     def _assemble_multi_class_output(self):
-        exprs = []
-        for idx in range(self.model.classes_.shape[0]):
-            exprs.append(self._assemble_single_output(idx))
+        exprs = [
+            self._assemble_single_output(idx)
+            for idx in range(self.model.classes_.shape[0])
+        ]
         return ast.VectorVal(exprs)
 
     def _get_single_coef(self, idx=0):
diff --git a/m2cgen/assemblers/utils.py b/m2cgen/assemblers/utils.py
index 5268a250..0be3b348 100644
--- a/m2cgen/assemblers/utils.py
+++ b/m2cgen/assemblers/utils.py
@@ -69,12 +69,13 @@ def _inner(current_expr, *rest_exprs):
 
 
 def to_1d_array(var):
-    return np.reshape(np.asarray(var), (np.size(var)))
+    return np.ravel(np.asarray(var))
 
 
 def to_2d_array(var):
-    if len(np.shape(var)) == 2:
-        x, y = var.shape
+    shape = np.shape(var)
+    if len(shape) == 2:
+        x, y = shape
     else:
         x, y = 1, np.size(var)
     return np.reshape(np.asarray(var), (x, y))
diff --git a/m2cgen/interpreters/code_generator.py b/m2cgen/interpreters/code_generator.py
index bf8d2f82..36401c55 100644
--- a/m2cgen/interpreters/code_generator.py
+++ b/m2cgen/interpreters/code_generator.py
@@ -53,9 +53,9 @@ def add_code_line(self, line):
 
     def add_code_lines(self, lines):
         if isinstance(lines, str):
-            lines = lines.split("\n")
-        for l in lines:
-            self.add_code_line(l)
+            lines = lines.strip().split("\n")
+        indent = " " * self._current_indent
+        self.code += indent + "\n{}".format(indent).join(lines) + "\n"
 
     def prepend_code_line(self, line):
         self.code = line + "\n" + self.code
@@ -63,8 +63,7 @@ def prepend_code_line(self, line):
     def prepend_code_lines(self, lines):
         if isinstance(lines, str):
             lines = lines.strip().split("\n")
-        for l in lines[::-1]:
-            self.prepend_code_line(l)
+        self.code = "\n".join(lines) + "\n" + self.code
 
     # Following methods simply compute expressions using templates without
     # changing result.