CODEWITHSUNDEEP
python-3.x
wwl
wwl

Reputation: 2065

Passing objects into yield statement in Python

I have a script that makes various read requests from a database (I have many methods similar to the read() method). I'd like to avoid repeating the code for creating and closing the cursor. So I am using a context manager and generator hand-in-hand. (In the below code, db_connection is a mysql connector object.)

from contextlib import contextmanager

class DatabaseReader():
  def __init__(self, db_connection):
    self.db_connection = db_connection
    self.num = num

  @contextmanager
  def connect(self):
    cursor = self.db_connection.cursor(buffered=True)
    try:
      yield
    finally:
      cursor.close()

  def read(self):
    with self.connect():
      cursor.execute("SELECT * FROM table;")

However, calling the 'read' method gives me an AttributeError as cursor is not defined. How can I "pass" the cursor object to the yield statement?

Upvotes: 0

Views: 846

Answers (1)

RomanPerekhrest
RomanPerekhrest

Reputation: 92854

In your custom implementation you need to "yield" a cursor object which will be bound in the as clause of the with statement:

3 fixes required in your code:

  • passing num to be assigned to self.num = num
  • yielding cursor
  • adding as keyword to with statement
from contextlib import contextmanager


class DatabaseReader():

    def __init__(self, db_connection, num):
        self.db_connection = db_connection
        self.num = num

    @contextmanager
    def connect(self):
        cursor = self.db_connection.cursor(buffered=True)
        try:
            yield cursor
        finally:
            cursor.close()

    def read(self):
        with self.connect() as cur:
            cur.execute("SELECT * FROM table;")

Upvotes: 3

Related Questions