[beam] branch master updated: [BEAM-13603] Fix bug in apache_beam.utils.Shared (#16437)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

altay pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new c67bed4  [BEAM-13603] Fix bug in apache_beam.utils.Shared (#16437)
c67bed4 is described below

commit c67bed4080734fdff9ae416483ce5591a98441f8
Author: ewianda <70...@users.noreply.github.com>
AuthorDate: Wed Jan 5 19:33:34 2022 -0500

    [BEAM-13603] Fix bug in apache_beam.utils.Shared (#16437)
    
    Co-authored-by: Ahmet Altay <aa...@gmail.com>
---
 sdks/python/apache_beam/utils/shared.py      |  1 +
 sdks/python/apache_beam/utils/shared_test.py | 11 +++++++++++
 2 files changed, 12 insertions(+)

diff --git a/sdks/python/apache_beam/utils/shared.py b/sdks/python/apache_beam/utils/shared.py
index 23622ef..295130b 100644
--- a/sdks/python/apache_beam/utils/shared.py
+++ b/sdks/python/apache_beam/utils/shared.py
@@ -147,6 +147,7 @@ class _SharedControlBlock(object):
         if result is None:
           return None
         self._ref = weakref.ref(result)
+        self._tag = tag
       else:
         result = self._ref()
     return result
diff --git a/sdks/python/apache_beam/utils/shared_test.py b/sdks/python/apache_beam/utils/shared_test.py
index caba2e5..48bd182 100644
--- a/sdks/python/apache_beam/utils/shared_test.py
+++ b/sdks/python/apache_beam/utils/shared_test.py
@@ -254,6 +254,17 @@ class SharedTest(unittest.TestCase):
     p2 = shared2.acquire(acquire_fn_2, tag='2')
     assert p2.get_name() == 'obj_2'
 
+  def testTagReturnsCached(self):
+    sequence = Sequence()
+    handle = shared.Shared()
+
+    f1 = handle.acquire(sequence.make_acquire_fn(), tag='1')
+    self.assertEqual('sequence1', f1.get_name())
+
+    # should return cached
+    f1 = handle.acquire(sequence.make_acquire_fn(), tag='1')
+    self.assertEqual('sequence1', f1.get_name())
+
 
 if __name__ == '__main__':
   unittest.main()