def __init__(self, out=sys.stdout): # Note that our 'database' could probably just be one dictionary. # However, since its possible that there are movies and actors # that share a name (e.g. "Ed Wood" <- A director, but you get # the idea), we break things up into 'actors' and 'films'. self._datastore = FilmGraph() self.out = out
def setUp(self): self.datastore = FilmGraph()
class FilmGraphTestCase(unittest.TestCase): def setUp(self): self.datastore = FilmGraph() def tearDown(self): pass def test_empty_datastore(self): expected = [] actual = self.datastore.get_shortest_path('Anyone') self.assertListEqual(expected, actual) def test_no_such_actor(self): self.datastore.add_link('John Crichton', 'Farscape') expected = [] actual = self.datastore.get_shortest_path('Aeryn Sun') self.assertListEqual(expected, actual) def test_no_such_target_actor(self): self.datastore.add_link('John Crichton', 'Farscape') expected = [] actual = self.datastore.get_shortest_path( 'John Crichton', 'Future Crichton' ) self.assertListEqual(expected, actual) def test_no_films_in_common(self): self.datastore.add_link('Superman', 'The Adventures of Lois and Clark') self.datastore.add_link('Clark Kent', 'Superman') expected = [] actual = self.datastore.get_shortest_path('Clark Kent', 'Superman') self.assertListEqual(expected, actual) def test_start_is_finish(self): self.datastore.add_link('Somebody to Love', 'Queen') expected = ['Somebody to Love'] actual = self.datastore.get_shortest_path( 'Somebody to Love', 'Somebody to Love' ) self.assertListEqual(expected, actual) def test_start_is_finish(self): self.datastore.add_link('Bill Murray', 'Ghostbusters') self.datastore.add_link('Bill Murray', 'SNL') self.datastore.add_link('John Belushi', 'SNL') self.datastore.add_link('John Belushi', 'Blues Brothers') self.datastore.add_link('Dan Aykroyd', 'Ghostbusters') self.datastore.add_link('Dan Aykroyd', 'Blues Brothers') expected = ['Bill Murray', 'Ghostbusters', 'Dan Aykroyd'] actual = self.datastore.get_shortest_path('Bill Murray', 'Dan Aykroyd') self.assertListEqual(expected, actual)
class Importer(object): """Handles the importing of film data (actors and film titles).""" def __init__(self, out=sys.stdout): # Note that our 'database' could probably just be one dictionary. # However, since its possible that there are movies and actors # that share a name (e.g. "Ed Wood" <- A director, but you get # the idea), we break things up into 'actors' and 'films'. self._datastore = FilmGraph() self.out = out @property def datastore(self): return self._datastore @datastore.setter def datastore(self, data): self._datastore = data def load_directory(self, directory): """Load film files (assumed to be JSON) from `directory`""" self.out.write('Attempting to load files from "{}"... '.format(directory)) try: files = os.listdir(directory) self.out.write('Success!\n') except: # We don't particularly care what the exception is. Just handle it. self.out.write( 'There was a problem accessing "{}"\n'.format(directory) ) return self # Iterate over all files in a folder for filename in files: path = os.path.join(directory, filename) self.load_file(path) return self def load_file(self, path): """Load film a film files (assumed to be JSON) from `path`""" self.out.write('\tLoading "{}"... '.format(path)) with open(path) as f: self.parse_file(f.read()) return self def parse_file(self, file_contents): """Parse contents of film file (assumed to be JSON) into our format.""" try: obj = json.loads(file_contents) except ValueError: # Skip this file. We only handle JSON right now. self.out.write('Skipping (Unsupported format)\n') return self title = obj.get('film', {}).get('name') if not title: self.out.write('Skipping (Missing title).\n') return self for actor in obj.get('cast'): name = actor.get('name') self._datastore.add_link(name, title) self.out.write('Success!\n') return self def stash(self, filename=None): """Pickle our datastore to stash it for later.""" if not filename: filename = settings.STASH_FILENAME pickle.dump(self.datastore, open(filename, 'wb')) return self def from_stash(self, filename=None): """Unpickle our datastore for further importing.""" if not filename: filename = settings.STASH_FILENAME self.datastore = pickle.load(open(filename, 'rb')) return self