Rewrite generate.py

Author: robertodr

chapter-06/recipe-03/cxx-example/generate.py

@@ -2,11 +2,12 @@
 Generates C++ vector of prime numbers up to max_number
 using sieve of Eratosthenes.
 """
+import pathlib
 import sys
 
 # for simplicity we do not verify argument list
 max_number = int(sys.argv[-2])
-output_file_name = sys.argv[-1]
+output_file_name = pathlib.Path(sys.argv[-1])
 
 numbers = range(2, max_number + 1)
 is_prime = {number: True for number in numbers}
@@ -19,12 +20,20 @@
             is_prime[current_position] = False
 
 primes = (number for number in numbers if is_prime[number])
-with open(output_file_name, 'w') as f:
-    f.write('#pragma once\n')
-    f.write('#include <vector>\n\n')
-    f.write('const std::size_t max_number = {0};\n'.format(max_number))
-    f.write('std::vector<int> & primes() {\n')
-    f.write('  static std::vector<int> primes;\n')
-    for number in primes:
-        f.write('  primes.push_back({0});\n'.format(number))
-    f.write('  return primes;\n}')
+code = """#pragma once
+
+#include <vector>
+
+const std::size_t max_number = {max_number};
+
+std::vector<int> & primes() {{
+  static std::vector<int> primes;
+
+{push_back}
+
+  return primes;
+}}
+"""
+push_back = '\n'.join(['  primes.push_back({:d});'.format(x) for x in primes])
+output_file_name.write_text(
+    code.format(max_number=max_number, push_back=push_back))