I have a Team class and a Driver class.
class Team:
def __init__(self, teamName, drivers=None):
self.teamName = teamName
if drivers is not None:
self.drivers = list(drivers)
else:
self.drivers = []
class Driver:
def __init__(self, name, country, number: str, podiums: int):
self.name = name
self.country = country
self.number = number
self.podiums = podiums
The Team class also has a function to add an instance of the Driver class to a particular team.
Each team has two drivers and the drivers can switch teams.
I can access the drivers that belong to a certain team but how do I access the team that a certain driver belongs to without having to manually add a team parameter in the Driver class?
There are many ways you could approach such a problem. Many represent a tradeoff between simplicity (and perhaps efficiency) on the one hand, and convenience (or "magic") on the other. Sometimes being straightforward and explicit really is best, even if it seems clunky.
That said... one way you might implement this is to give Team a class attribute that's a dictionary listing all team instances and their associated driver instances. I'm inferring from your question that a driver can only be on one team at a time, and that if possible you'd like that to be enforced as "automatically" as possible. I'm not sure if you want the code to enforce "no more than two drivers per team", perhaps throwing an error if you try to add a third, so I'll leave that aspect out for now.
class Team:
_membership = {} # Stores all teams and their members at the class level.
# The underscore indicates it's only intended to be used
# internally by the Team class.
def __init__(self, name, drivers=None):
self.name = name
Team._membership[self] = [] # Create an entry in the dictionary, with no drivers
if drivers is None:
drivers = []
for driver in drivers:
self.add_driver(driver) # Defined below
def __repr__(self):
"""Define how to represent this object in output."""
return f"<Team '{self.name}', drivers: {self.drivers}>"
We need to implement that add_driver() method. It will check if the driver's currently on a different team, and if so, remove it there:
def add_driver(self, driver):
old_team = Team.find_team_for_driver(driver) # Defined below
if old_team == self:
# If the driver is already on this team, we're done.
return
if old_team is not None:
# If it was on a different team, remove it from that team.
old_team.drivers.remove(driver)
# Lastly, add it to this team.
Team._membership[self].append(driver)
find_team_for_driver() isn't attached to a specific instance, but a method for general bookkeeping of the Team class, so it's a class method:
@classmethod
def find_team_for_driver(cls, driver):
for team, drivers in cls._membership.items():
if driver in drivers:
return team
# Return None if the driver was not found in a team.
return None
We can use a property decorator to make it easy to see what drivers a team has:
@property
def drivers(self):
return Team._membership[self]
@drivers.setter
def drivers(self, drivers):
Team._membership[self] = []
for driver in drivers:
self.add_driver(driver)
This means we can still use the simple syntax of team.drivers and the internals of the Team class will handle fetching it from the _membership dictionary. Note that while it's tempting to think you could add drivers to a team with team.drivers.append(new_driver), that won't know to remove the driver from their previous team, if any. Make sure to use team.add_driver(new_driver) instead.
(Note that team.drivers.remove(driver_to_remove) does work, but we could define a remove_driver() method if desired, just to keep things consistent.)
Okay, we're finally at the Driver class, which changes very little. (I've removed all attributes except name, just to make this example code a little simpler.) We can use a property to find out what team a driver is in and also to reassign it to a new team:
class Driver:
def __init__(self, name):
self.name = name
def __repr__(self):
return f"<Driver '{self.name}'>"
@property
def team(self):
return Team.find_team_for_driver(self)
@team.setter
def team(self, team):
team.add_driver(self)
Testing it out:
bob = Driver('Bob')
lisa = Driver('Lisa')
kaya = Driver('Kaya')
red = Team('Red Team', [bob, lisa])
print(red)
=> <Team 'Red Team', drivers: [<Driver 'Bob'>, <Driver 'Lisa'>]>
blue = Team('Blue Team')
blue.add_driver(kaya)
print(blue)
=> <Team 'Blue Team', drivers: [<Driver 'Kaya'>]>
Reassigning drivers:
lisa.team = blue
red.add_driver(kaya)
print(red)
=> <Team 'Red Team', drivers: [<Driver 'Bob'>, <Driver 'Kaya'>]>
print(blue)
=> <Team 'Blue Team', drivers: [<Driver 'Lisa'>]>
This kind of infrastructure can be handy for a certain size and complexity of project. For smaller, simpler things, you're probably better off explicitly defining and setting links between objects. (Note that all the iterating though Team._membership is slower than referencing explicit object attributes. This may not be enough of a difference to matter, in many cases.)
On the other end of the spectrum, as a project grows more complex, rather than implement more and more of this stuff, an existing system like Django starts to make more sense. Django's overall purpose is to be a web framework, but even if you're not doing anything web-related, its object-relational system implements all sorts of relations (one-to-one, one-to-many, many-to-many), with lots of convenient options like bidirectionality and enforcement of uniqueness. See this question if you get to a point where you'd like try that. (I believe SQLAlchemy implements something similar, but I've not tried it myself.)