Add auto-formatting and linting to our python code

.circleci/config.yml

@@ -6,7 +6,7 @@ orbs:
 parameters:
   image_suffix:
     type: string
-    default: '-v7e4468f'
+    default: '-v9a494cd'
   pg13_version:
     type: string
     default: '13.9'
@@ -157,8 +157,17 @@ jobs:
     steps:
       - checkout
       - run:
-          name: 'Check Style'
+          name: 'Check C Style'
           command: citus_indent --check
+      - run:
+          name: 'Check Python style'
+          command: black --check .
+      - run:
+          name: 'Check Python import order'
+          command: isort --check .
+      - run:
+          name: 'Check Python lints'
+          command: flake8 .
       - run:
           name: 'Fix whitespace'
           command: ci/editorconfig.sh && git diff --exit-code

.flake8

@@ -0,0 +1,7 @@
+[flake8]
+# E203 is ignored for black
+# E402 is ignored because of te way we do relative imports
+extend-ignore = E203, E402
+# black will truncate to 88 characters usually, but long string literals it
+# might keep. That's fine in most cases unless it gets really excessive.
+max-line-length = 150

.gitignore

@@ -51,3 +51,4 @@ lib*.pc
 
 # style related temporary outputs
 *.uncrustify
+.venv

Makefile

@@ -61,6 +61,9 @@ clean: clean-extension clean-pg_send_cancellation
 reindent:
 	${citus_abs_top_srcdir}/ci/fix_style.sh
 check-style:
+	black . --check --quiet
+	isort . --check --quiet
+	flake8
 	cd ${citus_abs_top_srcdir} && citus_indent --quiet --check
 .PHONY: reindent check-style
 

ci/fix_style.sh

@@ -9,6 +9,8 @@ cidir="${0%/*}"
 cd ${cidir}/..
 
 citus_indent . --quiet
+black . --quiet
+isort . --quiet
 ci/editorconfig.sh
 ci/remove_useless_declarations.sh
 ci/disallow_c_comments_in_migrations.sh

pyproject.toml

@@ -0,0 +1,5 @@
+[tool.isort]
+profile = 'black'
+
+[tool.black]
+include = '(src/test/regress/bin/diff-filter|\.pyi?|\.ipynb)$'

src/test/regress/Pipfile

@@ -10,6 +10,10 @@ docopt = "==0.6.2"
 cryptography = "==3.4.8"
 
 [dev-packages]
+black = "*"
+isort = "*"
+flake8 = "*"
+flake8-bugbear = "*"
 
 [requires]
 python_version = "3.9"

src/test/regress/bin/create_test.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python3
 
-import sys
-import random
 import os
+import random
+import sys
 
 if len(sys.argv) != 2:
     print(

src/test/regress/bin/diff-filter

@@ -5,69 +5,72 @@ diff-filter denormalizes diff output by having lines beginning with ' ' or '+'
 come from file2's unmodified version.
 """
 
-from sys import argv, stdin, stdout
 import re
+from sys import argv, stdin, stdout
 
 
 class FileScanner:
-	"""
-	FileScanner is an iterator over the lines of a file.
-	It can apply a rewrite rule which can be used to skip lines.
-	"""
-	def __init__(self, file, rewrite = lambda x:x):
-		self.file = file
-		self.line = 1
-		self.rewrite = rewrite
+    """
+    FileScanner is an iterator over the lines of a file.
+    It can apply a rewrite rule which can be used to skip lines.
+    """
+
+    def __init__(self, file, rewrite=lambda x: x):
+        self.file = file
+        self.line = 1
+        self.rewrite = rewrite
 
-	def __next__(self):
-		while True:
-			nextline = self.rewrite(next(self.file))
-			if nextline is not None:
-				self.line += 1
-				return nextline
+    def __next__(self):
+        while True:
+            nextline = self.rewrite(next(self.file))
+            if nextline is not None:
+                self.line += 1
+                return nextline
 
 
 def main():
-	# we only test //d rules, as we need to ignore those lines
-	regexregex = re.compile(r"^/(?P<rule>.*)/d$")
-	regexpipeline = []
-	for line in open(argv[1]):
-		line = line.strip()
-		if not line or line.startswith('#') or not line.endswith('d'):
-			continue
-		rule = regexregex.match(line)
-		if not rule:
-			raise 'Failed to parse regex rule: %s' % line
-		regexpipeline.append(re.compile(rule.group('rule')))
+    # we only test //d rules, as we need to ignore those lines
+    regexregex = re.compile(r"^/(?P<rule>.*)/d$")
+    regexpipeline = []
+    for line in open(argv[1]):
+        line = line.strip()
+        if not line or line.startswith("#") or not line.endswith("d"):
+            continue
+        rule = regexregex.match(line)
+        if not rule:
+            raise "Failed to parse regex rule: %s" % line
+        regexpipeline.append(re.compile(rule.group("rule")))
 
-	def sed(line):
-		if any(regex.search(line) for regex in regexpipeline):
-			return None
-		return line
+    def sed(line):
+        if any(regex.search(line) for regex in regexpipeline):
+            return None
+        return line
 
-	for line in stdin:
-		if line.startswith('+++ '):
-			tab = line.rindex('\t')
-			fname = line[4:tab]
-			file2 = FileScanner(open(fname.replace('.modified', ''), encoding='utf8'), sed)
-			stdout.write(line)
-		elif line.startswith('@@ '):
-			idx_start = line.index('+') + 1
-			idx_end = idx_start + 1
-			while line[idx_end].isdigit():
-				idx_end += 1
-			linenum = int(line[idx_start:idx_end])
-			while file2.line < linenum:
-				next(file2)
-			stdout.write(line)
-		elif line.startswith(' '):
-			stdout.write(' ')
-			stdout.write(next(file2))
-		elif line.startswith('+'):
-			stdout.write('+')
-			stdout.write(next(file2))
-		else:
-			stdout.write(line)
+    for line in stdin:
+        if line.startswith("+++ "):
+            tab = line.rindex("\t")
+            fname = line[4:tab]
+            file2 = FileScanner(
+                open(fname.replace(".modified", ""), encoding="utf8"), sed
+            )
+            stdout.write(line)
+        elif line.startswith("@@ "):
+            idx_start = line.index("+") + 1
+            idx_end = idx_start + 1
+            while line[idx_end].isdigit():
+                idx_end += 1
+            linenum = int(line[idx_start:idx_end])
+            while file2.line < linenum:
+                next(file2)
+            stdout.write(line)
+        elif line.startswith(" "):
+            stdout.write(" ")
+            stdout.write(next(file2))
+        elif line.startswith("+"):
+            stdout.write("+")
+            stdout.write(next(file2))
+        else:
+            stdout.write(line)
 
 
 main()

src/test/regress/citus_tests/arbitrary_configs/citus_arbitrary_configs.py

@@ -12,21 +12,23 @@
     --seed=<seed>                   random number seed
     --base                          whether to use the base sql schedule or not
 """
+import os
+import shutil
 import sys
-import os, shutil
 
 # https://stackoverflow.com/questions/14132789/relative-imports-for-the-billionth-time/14132912#14132912
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 
-import common
-import config as cfg
 import concurrent.futures
 import multiprocessing
-from docopt import docopt
-import time
 import random
+import time
 
+import common
+from docopt import docopt
+
+import config as cfg
 
 testResults = {}
 parallel_thread_amount = 1
@@ -115,7 +117,6 @@ def copy_copy_modified_binary(datadir):
 
 
 def copy_test_files(config):
-
     sql_dir_path = os.path.join(config.datadir, "sql")
     expected_dir_path = os.path.join(config.datadir, "expected")
 
@@ -132,18 +133,20 @@ def copy_test_files(config):
 
                 line = line[colon_index + 1 :].strip()
                 test_names = line.split(" ")
-                copy_test_files_with_names(test_names, sql_dir_path, expected_dir_path, config)
+                copy_test_files_with_names(
+                    test_names, sql_dir_path, expected_dir_path, config
+                )
 
 
 def copy_test_files_with_names(test_names, sql_dir_path, expected_dir_path, config):
     for test_name in test_names:
         # make empty files for the skipped tests
         if test_name in config.skip_tests:
             expected_sql_file = os.path.join(sql_dir_path, test_name + ".sql")
-            open(expected_sql_file, 'x').close()
+            open(expected_sql_file, "x").close()
 
             expected_out_file = os.path.join(expected_dir_path, test_name + ".out")
-            open(expected_out_file, 'x').close()
+            open(expected_out_file, "x").close()
 
             continue
 

src/test/regress/citus_tests/common.py

@@ -1,12 +1,12 @@
+import atexit
+import concurrent.futures
 import os
 import shutil
-import sys
 import subprocess
-import atexit
-import concurrent.futures
+import sys
 
 import utils
-from utils import USER, cd
+from utils import USER
 
 
 def initialize_temp_dir(temp_dir):
@@ -27,13 +27,11 @@ def initialize_temp_dir_if_not_exists(temp_dir):
 
 def parallel_run(function, items, *args, **kwargs):
     with concurrent.futures.ThreadPoolExecutor() as executor:
-        futures = [
-            executor.submit(function, item, *args, **kwargs)
-            for item in items
-        ]
+        futures = [executor.submit(function, item, *args, **kwargs) for item in items]
         for future in futures:
             future.result()
 
+
 def initialize_db_for_cluster(pg_path, rel_data_path, settings, node_names):
     subprocess.run(["mkdir", rel_data_path], check=True)
 
@@ -52,7 +50,7 @@ def initialize(node_name):
             "--encoding",
             "UTF8",
             "--locale",
-            "POSIX"
+            "POSIX",
         ]
         subprocess.run(command, check=True)
         add_settings(abs_data_path, settings)
@@ -72,11 +70,16 @@ def add_settings(abs_data_path, settings):
 
 def create_role(pg_path, node_ports, user_name):
     def create(port):
-        command = "SET citus.enable_ddl_propagation TO OFF; SELECT worker_create_or_alter_role('{}', 'CREATE ROLE {} WITH LOGIN CREATEROLE CREATEDB;', NULL)".format(
-            user_name, user_name
+        command = (
+            "SET citus.enable_ddl_propagation TO OFF;"
+            + "SELECT worker_create_or_alter_role('{}', 'CREATE ROLE {} WITH LOGIN CREATEROLE CREATEDB;', NULL)".format(
+                user_name, user_name
+            )
         )
         utils.psql(pg_path, port, command)
-        command = "SET citus.enable_ddl_propagation TO OFF; GRANT CREATE ON DATABASE postgres to {}".format(user_name)
+        command = "SET citus.enable_ddl_propagation TO OFF; GRANT CREATE ON DATABASE postgres to {}".format(
+            user_name
+        )
         utils.psql(pg_path, port, command)
 
     parallel_run(create, node_ports)
@@ -89,7 +92,9 @@ def coordinator_should_haveshards(pg_path, port):
     utils.psql(pg_path, port, command)
 
 
-def start_databases(pg_path, rel_data_path, node_name_to_ports, logfile_prefix, env_variables):
+def start_databases(
+    pg_path, rel_data_path, node_name_to_ports, logfile_prefix, env_variables
+):
     def start(node_name):
         abs_data_path = os.path.abspath(os.path.join(rel_data_path, node_name))
         node_port = node_name_to_ports[node_name]
@@ -248,7 +253,12 @@ def logfile_name(logfile_prefix, node_name):
 
 
 def stop_databases(
-    pg_path, rel_data_path, node_name_to_ports, logfile_prefix, no_output=False, parallel=True
+    pg_path,
+    rel_data_path,
+    node_name_to_ports,
+    logfile_prefix,
+    no_output=False,
+    parallel=True,
 ):
     def stop(node_name):
         abs_data_path = os.path.abspath(os.path.join(rel_data_path, node_name))
@@ -287,7 +297,9 @@ def initialize_citus_cluster(bindir, datadir, settings, config):
     initialize_db_for_cluster(
         bindir, datadir, settings, config.node_name_to_ports.keys()
     )
-    start_databases(bindir, datadir, config.node_name_to_ports, config.name, config.env_variables)
+    start_databases(
+        bindir, datadir, config.node_name_to_ports, config.name, config.env_variables
+    )
     create_citus_extension(bindir, config.node_name_to_ports.values())
     add_workers(bindir, config.worker_ports, config.coordinator_port())
     if not config.is_mx:
@@ -296,6 +308,7 @@ def initialize_citus_cluster(bindir, datadir, settings, config):
         add_coordinator_to_metadata(bindir, config.coordinator_port())
     config.setup_steps()
 
+
 def eprint(*args, **kwargs):
     """eprint prints to stderr"""