From d0e1e50135ec3118f3d4deb3a55f0e90e57abb2c Mon Sep 17 00:00:00 2001 From: DIGVIJAY <144053736+digvijay-y@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:01:51 +0000 Subject: [PATCH] Add get_by_comment method to NumberedObjectCollection and corresponding tests - Implemented a method to yield objects matching a comment string or regex pattern. - Added unit tests to validate functionality and error handling for the new method. --- montepy/numbered_object_collection.py | 39 +++++++++++++++++++++++++++ tests/test_numbered_collection.py | 15 +++++++++++ 2 files changed, 54 insertions(+) diff --git a/montepy/numbered_object_collection.py b/montepy/numbered_object_collection.py index 064742c3f..ab4dfb761 100644 --- a/montepy/numbered_object_collection.py +++ b/montepy/numbered_object_collection.py @@ -2,6 +2,7 @@ from __future__ import annotations from abc import ABC import itertools as it +import re import typing import weakref from numbers import Integral @@ -262,6 +263,44 @@ def check_number(self, number): f"Number {number} is already in use for the collection: {type(self).__name__} by {self[number]}" ) + def get_by_comment( + self, searcher: str | re.Pattern + ) -> typing.Generator[Numbered_MCNP_Object, None, None]: + """Yield all objects whose comments match the given text or pattern. + + The search is applied to each object's + :attr:`~montepy.mcnp_object.MCNP_Object.comments`, so both leading and + inline comments are considered. + + Parameters + ---------- + searcher : str or re.Pattern + A substring to search for, or a compiled regular expression. + + Returns + ------- + Generator[Numbered_MCNP_Object, None, None] + A generator of all matching objects. + + Raises + ------ + TypeError + if ``searcher`` is not a string or a regex compiled from a string. + """ + if isinstance(searcher, re.Pattern): + if not isinstance(searcher.pattern, str): + raise TypeError( + f"searcher must be a str, or a pattern compiled from a str. {searcher} given." + ) + elif not isinstance(searcher, str): + raise TypeError( + f"searcher must be a str, or a pattern compiled from a str. {searcher} given." + ) + + for obj in self: + if obj.comments.search(searcher): + yield obj + def _update_number(self, old_num, new_num, obj): """Updates the number associated with a specific object in the internal cache. diff --git a/tests/test_numbered_collection.py b/tests/test_numbered_collection.py index b4b67062b..b5106bebe 100644 --- a/tests/test_numbered_collection.py +++ b/tests/test_numbered_collection.py @@ -3,6 +3,7 @@ from hypothesis import given, settings, strategies as st import copy import itertools as it +import re import montepy import montepy.cells @@ -102,6 +103,20 @@ def test_check_number(self, cp_simple_problem): with pytest.raises(ValueError): cp_simple_problem.materials.check_number(-1) + def test_get_by_comment(self): + problem = montepy.read_input(os.path.join("tests", "inputs", "pin_cell.imcnp")) + + by_text = list(problem.cells.get_by_comment("uranium rod")) + assert by_text == [problem.cells[1]] + + by_regex = list(problem.cells.get_by_comment(re.compile(r"URANIUM", re.I))) + assert by_regex == [problem.cells[1]] + + assert list(problem.cells.get_by_comment("not present")) == [] + + with pytest.raises(TypeError): + list(problem.cells.get_by_comment(5)) + def test_update_number_not_in_cache(self): """Test that _update_number silently returns when object is not in cache.""" cells = montepy.Cells()