Add lint checks on examples/ and tests/, resolving or ignoring specific errors as as necessary. Also add pydoctest to the set of static analysis checks to github continuous integration.

PiperOrigin-RevId: 860166125

Author: Ryan McKenna

.github/workflows/ci.yml

@@ -25,17 +25,20 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install build wheel
-          pip install pylint pytype flake8 pylint-exit
+          pip install pylint pytype flake8 pylint-exit pydocstyle
           pip install -e .[dev]
           pip install -r docs/requirements.txt
           pip install crc32c
       - name: Run flake8
         run: flake8 jax_privacy tests examples
+      - name: Run pydocstyle
+        run: |
+          pydocstyle --convention=google --add-ignore=D101,D102,D103,D105,D202,D402 jax_privacy/
       - name: Run pylint
         run: |
-          pylint --ignore=auditing.py --rcfile=.pylintrc jax_privacy || pylint-exit -efail -wfail -cfail -rfail $?
-          # pylint --rcfile=.pylintrc examples || pylint-exit -efail -wfail -cfail -rfail $?
-          # pylint --rcfile=.pylintrc tests -d W0212,C0114 || pylint-exit -efail -wfail -cfail -rfail $?
+          pylint jax_privacy || pylint-exit -efail -wfail -cfail -rfail $?
+          pylint examples || pylint-exit -efail -wfail -cfail -rfail $?
+          pylint tests -d W0212,C0114 || pylint-exit -efail -wfail -cfail -rfail $?
         shell: bash
       - name: Run pytype
         run: |

examples/distributed_noise_generation.py

@@ -162,10 +162,10 @@ def run(pytree_like_model_params):
   t0 = time.time()
   compiled_run = run.lower(model_params).compile()
   t1 = time.time()
-  print('[BandMF] Compilation time: %.3f seconds' % (t1 - t0))
+  print(f'[BandMF] Compilation time: {t1-t0:.3f} seconds')
   state, noisy_grad = jax.block_until_ready(compiled_run(model_params))
   t2 = time.time()
-  print('[BandMF] Per-step run time: %.3f seconds' % ((t2 - t1) / steps))
+  print(f'[BandMF] Per-step run time: {(t2-t1)/steps:.3f} seconds')
 
   return state, noisy_grad
 

examples/keras_api_example.py

@@ -23,7 +23,8 @@
 from absl import app
 
 os.environ["KERAS_BACKEND"] = "jax"
-from jax_privacy import keras_api  # pylint: disable=g-import-not-at-top
+# pylint: disable=g-import-not-at-top,wrong-import-position
+from jax_privacy import keras_api
 import keras
 from keras import layers
 import numpy as np

jax_privacy/matrix_factorization/buffered_toeplitz.py

@@ -79,8 +79,7 @@ def check_float64_dtype(blt: 'BufferedToeplitz'):
 
 @dataclasses.dataclass(kw_only=True, frozen=True)
 class StreamingMatrixBuilder:
-  """Builder to convert a BLT to a StreamingMatrix.
-  """
+  """Builder to convert a BLT to a StreamingMatrix."""
 
   buf_decay: np.ndarray
   output_scale: np.ndarray

tests/clipping_test.py

@@ -31,11 +31,11 @@ def cartesian_product(**kwargs):
 
 PYTREE_STRUCTS = [
     jax.ShapeDtypeStruct(shape=(5, 5), dtype=jnp.float32),
-    dict(
-        a=jax.ShapeDtypeStruct(shape=(10,), dtype=jnp.float16),
-        b=jax.ShapeDtypeStruct(shape=(5, 5), dtype=jnp.float32),
-        c=jax.ShapeDtypeStruct(shape=(), dtype=jnp.bfloat16),
-    ),
+    {
+        'a': jax.ShapeDtypeStruct(shape=(10,), dtype=jnp.float16),
+        'b': jax.ShapeDtypeStruct(shape=(5, 5), dtype=jnp.float32),
+        'c': jax.ShapeDtypeStruct(shape=(), dtype=jnp.bfloat16),
+    },
 ]
 
 

tests/experimental/accounting_test.py

@@ -51,7 +51,7 @@
 PLD_ACCOUNTANT = functools.partial(
     dp_accounting.pld.PLDAccountant, value_discretization_interval=1e-3
 )
-RDP_ACCOUNTANT = dp_accounting.rdp.RdpAccountant
+RDP_ACCOUNTANT = dp_accounting.rdp.RdpAccountant  # pylint: disable=invalid-name
 
 
 def _make_test_case(event_fn, accountant):