diff --git a/recommendation_engine.py b/recommendation_engine.py index 19cbb5d..4e09bdd 100644 --- a/recommendation_engine.py +++ b/recommendation_engine.py @@ -28,6 +28,26 @@ class Engine: else: self.playlists = [] + def _patch_clips(self, clips): + inpoints = {} + for index, clip in enumerate(clips): + video_id = clip['id'].split('/')[0] + inpoints[video_id] = inpoints.get(video_id, []) + [{ + 'index': index, + 'position': clip['in'] + }] + for video_id in inpoints: + for i, inpoint in sorted( + inpoints[video_id], key=lambda inpoint: inpoint['position'] + ): + if i < len(inpoints[video_id]) - 1: + clips[inpoint['index']]['out'] = inpoints[video_id][i + 1] + else: + clips[inpoint['index']]['out'] self.pandora.get({ + 'id': id, keys: ['duration'] + })['duration'] + return clips + def _shift_clips(self, clips): index = random.randrange(len(clips)) return clips[index:] + clips[:index - 1] @@ -119,11 +139,11 @@ class Engine: 'id': storyline['id'], 'name': storyline['nodename'], 'tags': storyline['tags'], - 'clips': [{ + 'clips': self._patch_clips([{ 'id': clip['id'], 'in': clip['in'], 'out': clip['out'] - } for clip in clips if clip['value'] == storyline['name']] + } for clip in clips if clip['value'] == storyline['name']]) } for storyline in storylines] if playlist['clips']] with open(os.path.join(self.path, 'playlists.json'), 'w') as f: f.write(json.dumps(self.playlists, indent=4, sort_keys=True))