rth
rth

Reputation: 3106

wrapper class with __enter__/__exit__ in python3

I have a code that perfectly works in Python 2, but it doesn't work in Python 3.

There is an aggregator class data and a few classes to work with specific data formats.

class data():
   def __init__(self, file, format="x"):
      if format == "x":
         self.data = xdata(file)
      elif format == "y":
         self.data = ydata(file)
      # Redirect functions to the specific class
      self.__enter__       = self.data.__enter__
      self.__exit__        = self.data.__exit__
class xdata():
   def __init__(self, file):
      #do something
   def __enter__(self):
      return self
   def __exit__(self, exc_type, exc_value, traceback):
      #do something

class ydata():
   def __init__(self, file):
      #do something
   def __enter__(self):
      return self
   def __exit__(self,exc_type, exc_value, traceback):
      #do something

In python2 I was able to execute the following code without any errors,

with data("newfile.x") as sd:
    #do something with data

but python3 returns an error AttributeError: __enter__

Any ideas on how to fix it?

Upvotes: 1

Views: 61

Answers (1)

Patrick Haugh
Patrick Haugh

Reputation: 61042

__enter__ and __exit__ will be resolved as descriptors, which means that resolution bypasses the attributes of the class. You can provide your own descriptors for __enter__ and __exit__ using property:

class xdata():
   def __init__(self, file):
      self.attr = 'x'
   def __enter__(self):
      return self
   def __exit__(self, exc_type, exc_value, traceback):
      pass

class ydata():
   def __init__(self, file):
      self.attr = 'y'
   def __enter__(self):
      return self
   def __exit__(self, exc_type, exc_value, traceback):
      pass

class data():
   def __init__(self, file, format="x"):
      if format == "x":
         self.data = xdata(file)
      elif format == "y":
         self.data = ydata(file)
      # Redirect functions to the specific class
   @property
   def __enter__(self):
      return self.data.__enter__
   @property
   def __exit__(self):
      return self.data.__exit__

with data("", "x") as d:
   print(d.attr)  # x

with data("", "y") as d:
   print(d.attr)  # y

Upvotes: 1

Related Questions