X-Git-Url: https://git.pterodactylus.net/?a=blobdiff_plain;f=src%2Fmain%2Fjava%2Fnet%2Fpterodactylus%2Frhynodge%2Ffilters%2FEpisodeFilter.java;h=821c1f134c31e6ef8fc12c5b90673c01e3c45516;hb=7eaf9f52b3ef3eead4695ec9ef03611b71618078;hp=9318d848c39f1e6c21e116f9caa099caa82df224;hpb=d66a56f793fecb5b6fc8cc1f30fdf76edd56a277;p=rhynodge.git

diff --git a/src/main/java/net/pterodactylus/rhynodge/filters/EpisodeFilter.java b/src/main/java/net/pterodactylus/rhynodge/filters/EpisodeFilter.java
index 9318d84..821c1f1 100644
--- a/src/main/java/net/pterodactylus/rhynodge/filters/EpisodeFilter.java
+++ b/src/main/java/net/pterodactylus/rhynodge/filters/EpisodeFilter.java
@@ -19,10 +19,9 @@ package net.pterodactylus.rhynodge.filters;
 
 import static com.google.common.base.Optional.absent;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.FluentIterable.from;
 import static java.util.Arrays.asList;
 
-import java.util.HashMap;
-import java.util.Map;
 import java.util.Collection;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -35,7 +34,10 @@ import net.pterodactylus.rhynodge.states.FailedState;
 import net.pterodactylus.rhynodge.states.TorrentState;
 import net.pterodactylus.rhynodge.states.TorrentState.TorrentFile;
 
+import com.google.common.base.Function;
 import com.google.common.base.Optional;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimap;
 
 /**
  * {@link Filter} implementation that extracts {@link Episode} information from
@@ -63,19 +65,16 @@ public class EpisodeFilter implements Filter {
 		checkState(state instanceof TorrentState, "state is not a TorrentState but a %s!", state.getClass());
 
 		TorrentState torrentState = (TorrentState) state;
-		Map<Episode, Episode> episodes = new HashMap<Episode, Episode>();
+		final Multimap<Episode, TorrentFile> episodes = HashMultimap.create();
 		for (TorrentFile torrentFile : torrentState) {
 			Optional<Episode> episode = extractEpisode(torrentFile);
 			if (!episode.isPresent()) {
 				continue;
 			}
-			if (!episodes.containsKey(episode.get())) {
-				episodes.put(episode.get(), episode.get());
-			}
-			episodes.get(episode.get()).addTorrentFile(torrentFile);
+			episodes.put(episode.get(), torrentFile);
 		}
 
-		return new EpisodeState(episodes.values());
+		return new EpisodeState(from(episodes.keySet()).transform(episodeFiller(episodes)).toSet());
 	}
 
 	//
@@ -83,6 +82,27 @@ public class EpisodeFilter implements Filter {
 	//
 
 	/**
+	 * Returns a function that creates an {@link Episode} that contains all {@link
+	 * TorrentFile}s.
+	 *
+	 * @param episodeTorrents
+	 * 		A multimap mapping episodes to torrent files.
+	 * @return The function that performs the extraction of torrent files
+	 */
+	private static Function<Episode, Episode> episodeFiller(final Multimap<Episode, TorrentFile> episodeTorrents) {
+		return new Function<Episode, Episode>() {
+			@Override
+			public Episode apply(Episode episode) {
+				Episode completeEpisode = new Episode(episode.season(), episode.episode());
+				for (TorrentFile torrentFile : episodeTorrents.get(episode)) {
+					completeEpisode.addTorrentFile(torrentFile);
+				}
+				return completeEpisode;
+			}
+		};
+	}
+
+	/**
 	 * Extracts episode information from the given torrent file.
 	 *
 	 * @param torrentFile