Add name filter to attribute setter on `Segment` by jank324 · Pull Request #598 · desy-ml/cheetah

jank324 · 2025-12-04T12:45:38Z

Description

Adds a filter_name argument to the attribute setting convenience function of Segment, so it's possible to filter not only by type but also by name.

Also adds a convince getter function with similar behaviour that gets a single attribute from the filtered elements in the segment.

Motivation and Context

I have raised an issue to propose this change (required for new features and bug fixes)

We no longer have to write

for bpm_name in BPM_NAMES:
    getattr(segment, bpm_name).is_active = True

and can instead write

segment.set_attrs_on_every_element(
    filter_name=self.BPM_NAMES, is_active=True
)

Types of changes

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to change)
Documentation (update in the documentation)

Checklist

I have updated the changelog accordingly (required).
My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.
I have reformatted the code and checked that formatting passes (required).
I have have fixed all issues found by flake8 (required).
I have ensured that all pytest tests pass (required).
I have run pytest on a machine with a CUDA GPU and made sure all tests pass (required).
I have checked that the documentation builds (required).

Note: We are using a maximum length of 88 characters per line.

Copilot

Pull request overview

This PR adds a filter_name parameter to the set_attrs_on_every_element method in Segment, enabling filtering elements by name in addition to type. This provides a more convenient API for setting attributes on specific named elements.

Key changes:

Added filter_name parameter to set_attrs_on_every_element method
Updated filtering logic to support name-based filtering alongside type filtering
Ensured recursive calls pass the filter_name parameter to nested segments
Updated CHANGELOG to document the new feature

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
cheetah/accelerator/segment.py	Added `filter_name` parameter to `set_attrs_on_every_element`, updated filtering logic and recursive calls, and documented the new parameter
CHANGELOG.md	Added entry documenting the new `filter_name` feature

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-12-04T12:53:21Z

+        filter_name: str | None = None,
        is_recursive: bool = True,
        **kwargs: dict[str, Any],
    ) -> None:
        """
        Set attributes on every element of a specific type in the segment.

        :param filter_type: Type of the elements to set the attributes for.
+        :param filter_name: Names of the elements to set the attributes for.
        :param is_recursive: If `True`, the this method is applied to nested `Segment`s
            as well. If `False`, only the elements directly in the top-level `Segment`
            are considered.
        :param kwargs: Attributes to set and their values.
        """
        for element in self.elements:
-            if filter_type is None or isinstance(element, filter_type):
+            if (filter_type is None or isinstance(element, filter_type)) and (
+                filter_name is None or element.name == filter_name
+            ):
                for key, value in kwargs.items():
                    setattr(element, key, value)
            elif is_recursive and isinstance(element, Segment):
                element.set_attrs_on_every_element(
-                    filter_type=filter_type, is_recursive=True, **kwargs
+                    filter_type=filter_type,
+                    filter_name=filter_name,
+                    is_recursive=True,
+                    **kwargs,
                )


The new filter_name parameter lacks test coverage. Since the existing function set_attrs_on_every_element has test coverage in tests/test_segment.py (see test_attr_setting_by_element_type_convenience_method), the new functionality should also be tested.

Consider adding a test that:

Creates a segment with elements having specific names

Calls set_attrs_on_every_element with the filter_name parameter

Verifies that only elements matching the specified name have their attributes modified

Tests the recursive behavior with nested segments

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 6 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-12-04T13:05:20Z

+    def get_attr_from_every_element(
+        self,
+        filter_type: type[Element] | tuple[type[Element]] | None = None,
+        filter_name: str | None = None,


Inconsistent type annotation for filter_name parameter. In get_attr_from_every_element, it's typed as str | None, but in set_attrs_on_every_element (line 622), it's typed as str | tuple[str] | None. These should be consistent. The setter allows a tuple to support multiple names, and the getter should do the same for consistency.

@copilot open a new pull request to apply changes based on this feedback

Copilot · 2025-12-04T13:05:21Z

+    def get_attr_from_every_element(
+        self,
+        filter_type: type[Element] | tuple[type[Element]] | None = None,
+        filter_name: str | None = None,
+        is_recursive: bool = True,
+    ) -> list[Any]:
+        """
+        Get an attribute from every element type in the segment filtered by type and/or
+        name.
+        :param filter_type: Type of the elements to get the attribute from.
+        :param filter_name: Name of the elements to get the attribute from.
+        :param is_recursive: If `True`, this method is applied to nested `Segment`s as
+            well. If `False`, only the elements directly in the top-level `Segment` are
+            considered.
+        :return: List of attributes from the filtered elements.
+        """
+        attrs = []
+        for element in self.elements:
+            if (filter_type is None or isinstance(element, filter_type)) and (
+                filter_name is None or element.name in filter_name
+            ):
+                attrs.append(element)
+            elif is_recursive and isinstance(element, Segment):
+                attrs.extend(
+                    element.get_attr_from_every_element(
+                        filter_type=filter_type,
+                        filter_name=filter_name,
+                        is_recursive=True,
+                    )
+                )
+        return attrs


The new filter_name parameter and get_attr_from_every_element method lack test coverage. There are existing tests for set_attrs_on_every_element with filter_type (line 57-86 in test_segment.py), but no tests for the new filtering by name functionality or the new getter method.

Copilot · 2025-12-04T13:05:22Z

+            if (filter_type is None or isinstance(element, filter_type)) and (
+                filter_name is None or element.name in filter_name
+            ):
                for key, value in kwargs.items():
                    setattr(element, key, value)
            elif is_recursive and isinstance(element, Segment):
                element.set_attrs_on_every_element(
-                    filter_type=filter_type, is_recursive=True, **kwargs
+                    filter_type=filter_type,
+                    filter_name=filter_name,
+                    is_recursive=True,
+                    **kwargs,
                )


Missing test coverage for the new filter_name parameter. While tests exist for filtering by type, the new name filtering functionality is not tested.

@copilot open a new pull request to apply changes based on this feedback

Copilot · 2025-12-04T13:30:35Z

@jank324 I've opened a new pull request, #599, to work on those changes. Once the pull request is ready, I'll request review from you.

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Copilot · 2025-12-04T13:35:11Z

@jank324 I've opened a new pull request, #600, to work on those changes. Once the pull request is ready, I'll request review from you.

jank324 · 2026-03-05T09:50:26Z


 ### 🚀 Features

+- Add `filter_name` argument to the attribute setting convenience function of Segment, so it's possible to filter not only by type but also by name. (see #598) (@jank324)


Needs to be moved to new version

Add name filter to attribute setter on Segment

541a74c

jank324 self-assigned this Dec 4, 2025

jank324 added the enhancement New feature or request label Dec 4, 2025

Add changelog entry

5cd77be

jank324 marked this pull request as ready for review December 4, 2025 12:49

jank324 requested review from Hespe and Copilot December 4, 2025 12:50

Copilot started reviewing on behalf of jank324 December 4, 2025 12:50 View session

Copilot finished reviewing on behalf of jank324 December 4, 2025 12:52

Copilot AI reviewed Dec 4, 2025

View reviewed changes

jank324 added 4 commits December 4, 2025 13:56

Add a convenience attribute getter method

f1c5d71

Add a second changelog entry

8cc12c3

Fix name filter issue pointed out by Copilot

786184d

Fix docstring typo pointed out by Copilot

aae4436

jank324 requested a review from Copilot December 4, 2025 13:01

Copilot started reviewing on behalf of jank324 December 4, 2025 13:02 View session

Copilot finished reviewing on behalf of jank324 December 4, 2025 13:04

Copilot AI reviewed Dec 4, 2025

View reviewed changes

Copilot AI mentioned this pull request Dec 4, 2025

Fix filter_name type consistency between getter and setter methods #599

Draft

14 tasks

jank324 and others added 3 commits December 4, 2025 14:31

Apply Copilot suggestion

713be0c

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Fix typo highlighted by Copilot

3f5a23e

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Update cheetah/accelerator/segment.py

0834629

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Copilot AI mentioned this pull request Dec 4, 2025

Add test coverage for filter_name parameter in Segment methods #600

Draft

14 tasks

jank324 commented Mar 5, 2026

View reviewed changes


		### 🚀 Features

		- Add `filter_name` argument to the attribute setting convenience function of Segment, so it's possible to filter not only by type but also by name. (see #598) (@jank324)

Conversation

jank324 commented Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Types of changes

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

jank324 Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

jank324 Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Copilot AI commented Dec 4, 2025

Uh oh!

Copilot AI commented Dec 4, 2025

Uh oh!

jank324 Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jank324 commented Dec 4, 2025 •

edited

Loading