apply the C4 rule (flake8-comprehensions)

data/json/json2yaml.py

@@ -19,7 +19,7 @@ def _main():
     )
 
     # get all json files in dir
-    jsons = [p for p in Path.cwd().glob("*.json")]
+    jsons = list(Path.cwd().glob("*.json"))
     # use the newest as autosuggestion
     jsons.sort(key=lambda x: x.stat().st_mtime, reverse=True)
     jfile = jsons[0]

deepmd/descriptor/se_a.py

@@ -1005,10 +1005,7 @@ def _filter(
         outputs_size = [1] + self.filter_neuron
         outputs_size_2 = self.n_axis_neuron
         all_excluded = all(
-            [
-                (type_input, type_i) in self.exclude_types
-                for type_i in range(self.ntypes)
-            ]
+            (type_input, type_i) in self.exclude_types for type_i in range(self.ntypes)
         )
         if all_excluded:
             # all types are excluded so result and qmat should be zeros

deepmd/fit/dipole.py

@@ -75,7 +75,7 @@ def __init__(
         self.resnet_dt = resnet_dt
         self.sel_type = sel_type
         if self.sel_type is None:
-            self.sel_type = [ii for ii in range(self.ntypes)]
+            self.sel_type = list(range(self.ntypes))
         self.sel_mask = np.array(
             [ii in self.sel_type for ii in range(self.ntypes)], dtype=bool
         )

deepmd/fit/ener.py

@@ -283,11 +283,11 @@ def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False):
             # In this situation, we directly use these assigned energies instead of computing stats.
             # This will make the loss decrease quickly
             assigned_atom_ener = np.array(
-                list(ee for ee in self.atom_ener_v if ee is not None)
-            )
-            assigned_ener_idx = list(
-                (ii for ii, ee in enumerate(self.atom_ener_v) if ee is not None)
+                [ee for ee in self.atom_ener_v if ee is not None]
             )
+            assigned_ener_idx = [
+                ii for ii, ee in enumerate(self.atom_ener_v) if ee is not None
+            ]
             # np.dot out size: nframe
             sys_ener -= np.dot(sys_tynatom[:, assigned_ener_idx], assigned_atom_ener)
             sys_tynatom[:, assigned_ener_idx] = 0.0

deepmd/fit/polar.py

@@ -98,7 +98,7 @@ def __init__(
         self.fitting_activation_fn = get_activation_func(activation_function)
         self.fitting_precision = get_precision(precision)
         if self.sel_type is None:
-            self.sel_type = [ii for ii in range(self.ntypes)]
+            self.sel_type = list(range(self.ntypes))
         self.sel_mask = np.array(
             [ii in self.sel_type for ii in range(self.ntypes)], dtype=bool
         )

deepmd/infer/deep_dos.py

@@ -63,29 +63,27 @@ def __init__(
         # add these tensors on top of what is defined by DeepTensor Class
         # use this in favor of dict update to move attribute from class to
         # instance namespace
-        self.tensors = dict(
-            {
-                # descrpt attrs
-                "t_ntypes": "descrpt_attr/ntypes:0",
-                "t_rcut": "descrpt_attr/rcut:0",
-                # fitting attrs
-                "t_dfparam": "fitting_attr/dfparam:0",
-                "t_daparam": "fitting_attr/daparam:0",
-                "t_numb_dos": "fitting_attr/numb_dos:0",
-                # model attrs
-                "t_tmap": "model_attr/tmap:0",
-                # inputs
-                "t_coord": "t_coord:0",
-                "t_type": "t_type:0",
-                "t_natoms": "t_natoms:0",
-                "t_box": "t_box:0",
-                "t_mesh": "t_mesh:0",
-                # add output tensors
-                "t_dos": "o_dos:0",
-                "t_atom_dos": "o_atom_dos:0",
-                "t_descriptor": "o_descriptor:0",
-            },
-        )
+        self.tensors = {
+            # descrpt attrs
+            "t_ntypes": "descrpt_attr/ntypes:0",
+            "t_rcut": "descrpt_attr/rcut:0",
+            # fitting attrs
+            "t_dfparam": "fitting_attr/dfparam:0",
+            "t_daparam": "fitting_attr/daparam:0",
+            "t_numb_dos": "fitting_attr/numb_dos:0",
+            # model attrs
+            "t_tmap": "model_attr/tmap:0",
+            # inputs
+            "t_coord": "t_coord:0",
+            "t_type": "t_type:0",
+            "t_natoms": "t_natoms:0",
+            "t_box": "t_box:0",
+            "t_mesh": "t_mesh:0",
+            # add output tensors
+            "t_dos": "o_dos:0",
+            "t_atom_dos": "o_atom_dos:0",
+            "t_descriptor": "o_descriptor:0",
+        }
         DeepEval.__init__(
             self,
             model_file,

deepmd/infer/deep_pot.py

@@ -78,31 +78,29 @@ def __init__(
         # add these tensors on top of what is defined by DeepTensor Class
         # use this in favor of dict update to move attribute from class to
         # instance namespace
-        self.tensors = dict(
-            {
-                # descrpt attrs
-                "t_ntypes": "descrpt_attr/ntypes:0",
-                "t_rcut": "descrpt_attr/rcut:0",
-                # fitting attrs
-                "t_dfparam": "fitting_attr/dfparam:0",
-                "t_daparam": "fitting_attr/daparam:0",
-                # model attrs
-                "t_tmap": "model_attr/tmap:0",
-                # inputs
-                "t_coord": "t_coord:0",
-                "t_type": "t_type:0",
-                "t_natoms": "t_natoms:0",
-                "t_box": "t_box:0",
-                "t_mesh": "t_mesh:0",
-                # add output tensors
-                "t_energy": "o_energy:0",
-                "t_force": "o_force:0",
-                "t_virial": "o_virial:0",
-                "t_ae": "o_atom_energy:0",
-                "t_av": "o_atom_virial:0",
-                "t_descriptor": "o_descriptor:0",
-            },
-        )
+        self.tensors = {
+            # descrpt attrs
+            "t_ntypes": "descrpt_attr/ntypes:0",
+            "t_rcut": "descrpt_attr/rcut:0",
+            # fitting attrs
+            "t_dfparam": "fitting_attr/dfparam:0",
+            "t_daparam": "fitting_attr/daparam:0",
+            # model attrs
+            "t_tmap": "model_attr/tmap:0",
+            # inputs
+            "t_coord": "t_coord:0",
+            "t_type": "t_type:0",
+            "t_natoms": "t_natoms:0",
+            "t_box": "t_box:0",
+            "t_mesh": "t_mesh:0",
+            # add output tensors
+            "t_energy": "o_energy:0",
+            "t_force": "o_force:0",
+            "t_virial": "o_virial:0",
+            "t_ae": "o_atom_energy:0",
+            "t_av": "o_atom_virial:0",
+            "t_descriptor": "o_descriptor:0",
+        }
         DeepEval.__init__(
             self,
             model_file,

deepmd/train/trainer.py

@@ -1200,4 +1200,4 @@ def get_data_dict(self, batch_list: List[np.ndarray]) -> Dict[str, np.ndarray]:
         Dict[str, np.ndarray]
             The dict of the loaded data.
         """
-        return {kk: vv for kk, vv in zip(self.data_keys, batch_list)}
+        return dict(zip(self.data_keys, batch_list))

deepmd/utils/multi_init.py

@@ -66,9 +66,9 @@ def replace_model_params_with_frz_multi_model(
 
     # Change model configurations
     pretrained_fitting_keys = sorted(
-        list(pretrained_jdata["model"]["fitting_net_dict"].keys())
+        pretrained_jdata["model"]["fitting_net_dict"].keys()
     )
-    cur_fitting_keys = sorted(list(jdata["model"]["fitting_net_dict"].keys()))
+    cur_fitting_keys = sorted(jdata["model"]["fitting_net_dict"].keys())
     newly_added_fittings = set(cur_fitting_keys) - set(pretrained_fitting_keys)
     reused_fittings = set(cur_fitting_keys) - newly_added_fittings
     log.info("Change the model configurations according to the pretrained one...")

deepmd/utils/path.py

@@ -173,7 +173,7 @@ def glob(self, pattern: str) -> List["DPPath"]:
         """
         # currently DPOSPath will only derivative DPOSPath
         # TODO: discuss if we want to mix DPOSPath and DPH5Path?
-        return list([type(self)(p) for p in self.path.glob(pattern)])
+        return [type(self)(p) for p in self.path.glob(pattern)]
 
     def rglob(self, pattern: str) -> List["DPPath"]:
         """This is like calling :meth:`DPPath.glob()` with `**/` added in front
@@ -189,7 +189,7 @@ def rglob(self, pattern: str) -> List["DPPath"]:
         List[DPPath]
             list of paths
         """
-        return list([type(self)(p) for p in self.path.rglob(pattern)])
+        return [type(self)(p) for p in self.path.rglob(pattern)]
 
     def is_file(self) -> bool:
         """Check if self is file."""
@@ -290,12 +290,10 @@ def glob(self, pattern: str) -> List["DPPath"]:
         """
         # got paths starts with current path first, which is faster
         subpaths = [ii for ii in self._keys if ii.startswith(self.name)]
-        return list(
-            [
-                type(self)(f"{self.root_path}#{pp}")
-                for pp in globfilter(subpaths, self._connect_path(pattern))
-            ]
-        )
+        return [
+            type(self)(f"{self.root_path}#{pp}")
+            for pp in globfilter(subpaths, self._connect_path(pattern))
+        ]
 
     def rglob(self, pattern: str) -> List["DPPath"]:
         """This is like calling :meth:`DPPath.glob()` with `**/` added in front

deepmd/utils/tabulate.py

@@ -705,9 +705,7 @@ def _all_excluded(self, ii: int) -> bool:
         bool
             if type ii excluds all types
         """
-        return all(
-            [(ii, type_i) in self.exclude_types for type_i in range(self.ntypes)]
-        )
+        return all((ii, type_i) in self.exclude_types for type_i in range(self.ntypes))
 
     def _get_table_size(self):
         table_size = 0

deepmd/utils/type_embed.py

@@ -128,7 +128,7 @@ def build(
         embedded_types
             The computational graph for embedded types
         """
-        types = tf.convert_to_tensor([ii for ii in range(ntypes)], dtype=tf.int32)
+        types = tf.convert_to_tensor(list(range(ntypes)), dtype=tf.int32)
         ebd_type = tf.cast(
             tf.one_hot(tf.cast(types, dtype=tf.int32), int(ntypes)),
             self.filter_precision,

pyproject.toml

@@ -129,6 +129,7 @@ select = [
     "F", # pyflakes
     "D", # pydocstyle
     "UP", # pyupgrade
+    "C4", # flake8-comprehensions
 ]
 ignore = [
     "E501", # line too long

source/install/build_tf.py

@@ -351,16 +351,14 @@ def include_patterns(*include_patterns):
     """
 
     def _ignore_patterns(path, names):
-        keep = set(
-            name for pattern in include_patterns for name in filter(names, pattern)
-        )
-        removed_dir = any([x.startswith("_") for x in path.split(os.path.sep)])
-        ignore = set(
+        keep = {name for pattern in include_patterns for name in filter(names, pattern)}
+        removed_dir = any(x.startswith("_") for x in path.split(os.path.sep))
+        ignore = {
             name
             for name in names
             if (name not in keep or removed_dir)
             and not os.path.isdir(os.path.join(path, name))
-        )
+        }
         return ignore
 
     return _ignore_patterns