kirpit kirpit - 5 months ago 13
Python Question

Pythonic way to yield / return main item within a sub loop

I have two different objects. One of them wraps N number of the other type objects within a list or tuple attribute. Let's say students in classrooms:

class Student:
def __init__(self, name):
self.name = name

class ClassRoom:
def __init__(self, students):
self.students = students


Of course, we have heaps of Student as well as ClassRoom instances:

john, sam = Student('John'), Student('Sam')
patrick, michael, bill = Student('Patrick'), Student('Michael'), Student('Bill')
klass1 = ClassRoom([john, sam])
klass2 = ClassRoom([patrick, michael, bill])


Just consider that each student is unique by their name and you cannot reach to a student's classroom by reference as such:

sam.get_classroom() # Student class doesn't support this :(


And we have a helper function that does the job:

def get_classroom_by_student(klasses, student_name):
for klass in klasses:
for student in klass.students:
if student.name==student_name:
return klass
# Or yield if a student belongs to more than one class

sams_class = get_classroom_by_student([klass1, klass2], 'Sam')
bills_class = get_classroom_by_student([klass1, klass2], 'Bill')


Since "Flat is better than nested", how do I create an efficient generator or is there some pythonic way to implement this helper function?

Answer

How about this:

class Student:
    def __init__(self, name):
        self.name = name

class ClassRoom:
    def __init__(self, students):
        self.students = students

john, sam = Student('John'), Student('Sam')
patrick, michael, bill = Student('Patrick'), Student('Michael'), Student('Bill')

klass1 = ClassRoom([john, sam])
klass2 = ClassRoom([patrick, michael, bill])


def where_is(student, klasses):
    return next((x for x in klasses if student in x.students), None)

assert klass1 is where_is(john, [klass1, klass2])
assert klass2 is where_is(patrick, [klass1, klass2])

nobody = Student('foo')
assert None is where_is(nobody, [klass1, klass2])

For the yielding version just omit next and return the generator:

def where_is(student, klasses):
    return (x for x in klasses if student in x.students)

for klass in where_is(john, [klass1, klass2]):
    print klass