def identify_fields()

in robogym/randomization/sim.py [0:0]

• 47 lines of code
• 16 McCabe index (conditional complexity)

    def identify_fields(self, geom_prefix, body_prefix, dof_jnt_prefix, jnt_prefix):
        if geom_prefix is not None:
            assert self._field_name.startswith("geom_")
            geom_names = [
                name
                for name in self.sim.model.geom_names
                if has_prefixes(name, geom_prefix)
            ]
            ids = np.array(
                sorted([self.sim.model.geom_name2id(name) for name in geom_names])
            )
        elif body_prefix is not None:
            assert self._field_name.startswith("body_")
            body_names = [
                name
                for name in self.sim.model.body_names
                if has_prefixes(name, body_prefix)
            ]
            ids = np.array(
                sorted([self.sim.model.body_name2id(name) for name in body_names])
            )
        elif dof_jnt_prefix is not None:

            def has_prefix(jnt_id):
                return has_prefixes(
                    self.sim.model.joint_id2name(jnt_id), dof_jnt_prefix
                )

            assert self._field_name.startswith("dof_")
            ids = [
                idx
                for idx, jnt_id in enumerate(self.sim.model.dof_jntid)
                if has_prefix(jnt_id)
            ]
        elif jnt_prefix is not None:

            def has_prefix(jnt_id):
                return has_prefixes(self.sim.model.joint_id2name(jnt_id), jnt_prefix)

            assert self._field_name.startswith("jnt_")
            ids = [
                idx for idx in range(len(self.sim.model.jnt_type)) if has_prefix(idx)
            ]
        else:
            ids = None

        if ids is not None:
            ids = np.array(sorted(ids))
            assert len(ids) > 0, "no IDs matched for {}".format(self._field_name)
        else:
            ids = None

        return ids