simpleperf: expose SetSampleFilter in python report lib.

Bug: 211814099
Test: run test.py --only-host-test
Change-Id: I71fc59c758f795d7644a9091fe7fdcd51c7d22c3

2 files changed, 60 insertions, 0 deletions

diff --git a/simpleperf/scripts/simpleperf_report_lib.py b/simpleperf/scripts/simpleperf_report_lib.py
index e26a03ff..5c2893c1 100644
--- a/simpleperf/scripts/simpleperf_report_lib.py
+++ b/simpleperf/scripts/simpleperf_report_lib.py
@@ -261,6 +261,8 @@ class ReportLib(object):
         self._GetSupportedTraceOffCpuModesFunc.restype = ct.c_char_p
         self._SetTraceOffCpuModeFunc = self._lib.SetTraceOffCpuMode
         self._SetTraceOffCpuModeFunc.restype = ct.c_bool
+        self._SetSampleFilterFunc = self._lib.SetSampleFilter
+        self._SetSampleFilterFunc.restype = ct.c_bool
         self._GetNextSampleFunc = self._lib.GetNextSample
         self._GetNextSampleFunc.restype = ct.POINTER(SampleStruct)
         self._GetEventOfCurrentSampleFunc = self._lib.GetEventOfCurrentSample
@@ -360,6 +362,28 @@ class ReportLib(object):
         res: bool = self._SetTraceOffCpuModeFunc(self.getInstance(), _char_pt(mode))
         _check(res, f'Failed to call SetTraceOffCpuMode({mode})')
 
+    def SetSampleFilter(self, filter: str):
+        """ Set options used to filter samples. Available options are:
+            --exclude-pid pid1,pid2,...   Exclude samples for selected processes.
+            --exclude-tid tid1,tid2,...   Exclude samples for selected threads.
+            --exclude-process-name process_name_regex   Exclude samples for processes with name
+                                                        containing the regular expression.
+            --exclude-thread-name thread_name_regex     Exclude samples for threads with name
+                                                        containing the regular expression.
+            --include-pid pid1,pid2,...   Include samples for selected processes.
+            --include-tid tid1,tid2,...   Include samples for selected threads.
+            --include-process-name process_name_regex   Include samples for processes with name
+                                                        containing the regular expression.
+            --include-thread-name thread_name_regex     Include samples for threads with name
+                                                        containing the regular expression.
+            --filter-file <file>          Use filter file to filter samples based on timestamps. The
+                                          file format is in doc/sampler_filter.md.
+
+            The filter argument should be a concatenation of options.
+        """
+        res: bool = self._SetSampleFilterFunc(self.getInstance(), _char_pt(filter))
+        _check(res, f'Failed to call SetSampleFilter({filter})')
+
     def GetNextSample(self) -> Optional[SampleStruct]:
         """ Return the next sample. If no more samples, return None. """
         psample = self._GetNextSampleFunc(self.getInstance())
diff --git a/simpleperf/scripts/test/report_lib_test.py b/simpleperf/scripts/test/report_lib_test.py
index d4c3f091..61899459 100644
--- a/simpleperf/scripts/test/report_lib_test.py
+++ b/simpleperf/scripts/test/report_lib_test.py
@@ -14,6 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import tempfile
+from typing import Set
+
 from simpleperf_report_lib import ReportLib
 from . test_utils import TestBase, TestHelper
 
@@ -268,3 +271,36 @@ class TestReportLib(TestBase):
         self.assertEqual(self.report_lib.GetSupportedTraceOffCpuModes(), [])
         with self.assertRaises(RuntimeError):
             self.report_lib.SetTraceOffCpuMode('on-cpu')
+
+    def test_set_sample_filter(self):
+        """ Test using ReportLib.SetSampleFilter(). """
+        def get_threads_for_filter(filter: str) -> Set[int]:
+            self.report_lib.Close()
+            self.report_lib = ReportLib()
+            self.report_lib.SetRecordFile(TestHelper.testdata_path('perf_display_bitmaps.data'))
+            self.report_lib.SetSampleFilter(filter)
+            threads = set()
+            while self.report_lib.GetNextSample():
+                sample = self.report_lib.GetCurrentSample()
+                threads.add(sample.tid)
+            return threads
+
+        self.assertNotIn(31850, get_threads_for_filter('--exclude-pid 31850'))
+        self.assertIn(31850, get_threads_for_filter('--include-pid 31850'))
+        self.assertNotIn(31881, get_threads_for_filter('--exclude-tid 31881'))
+        self.assertIn(31881, get_threads_for_filter('--include-tid 31881'))
+        self.assertNotIn(31881, get_threads_for_filter(
+            '--exclude-process-name com.example.android.displayingbitmaps'))
+        self.assertIn(31881, get_threads_for_filter(
+            '--include-process-name com.example.android.displayingbitmaps'))
+        self.assertNotIn(31850, get_threads_for_filter(
+            '--exclude-thread-name com.example.android.displayingbitmaps'))
+        self.assertIn(31850, get_threads_for_filter(
+            '--include-thread-name com.example.android.displayingbitmaps'))
+
+        with tempfile.NamedTemporaryFile('w') as filter_file:
+            filter_file.write('GLOBAL_BEGIN 684943449406175\nGLOBAL_END 684943449406176')
+            filter_file.flush()
+            threads = get_threads_for_filter('--filter-file ' + filter_file.name)
+            self.assertIn(31881, threads)
+            self.assertNotIn(31850, threads)