Allow multiple targets to be enabled/disabled

3 files changed, 273 insertions, 28 deletions

diff --git a/src/main/TweetShell.java b/src/main/TweetShell.java
index 5e5d074..fa1fb49 100644
--- a/src/main/TweetShell.java
+++ b/src/main/TweetShell.java
@@ -5,17 +5,21 @@ import io.OAuthRequester;
 import io.StreamImpl;
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Scanner;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import mining.Stream;
 import mining.TwitterApi;
 import org.json.JSONException;
 import org.json.JSONObject;
+import provider.CompositeResultListener;
 import provider.ExceptionListener;
 import provider.ResultListener;
+import support.ClassEnabledTracker;
 import utils.Configuration;
 
 /**
@@ -42,7 +46,13 @@ public class TweetShell implements TwitterApi.PinSupplier {
             stream_cached = new StreamImpl(requester);
             StreamImpl streamObserver = (StreamImpl) stream_cached;
             streamObserver.setExceptionListener(handler);
-            streamObserver.setResultListener(new StreamHandler());
+
+            CompositeResultListener listeners = new CompositeResultListener();
+            streamObserver.setResultListener(listeners);
+            // by default, store something that counts responses
+            listeners.register(new TweetCounter());
+            // and something that prints tweets to console.
+            listeners.register(handler);
         }
         return stream_cached;
     }
@@ -155,7 +165,9 @@ public class TweetShell implements TwitterApi.PinSupplier {
         close("Close the stream"),
         exit("Returns to shell"),
         help("Get help"),
-        target("Set output target: {file, shell}.", 1);
+        target("Set output target: one or more of {file, shell} with optional."
+                + "'-' (disable) or '+' (enable) prefix. Without prefix, "
+                + "only that target is enabled.");
 
         private final String description;
         private final int paramCount;
@@ -258,37 +270,119 @@ public class TweetShell implements TwitterApi.PinSupplier {
                 throw new NoSuchElementException();
             case target:
                 StreamImpl stream = (StreamImpl) getStream();
-                ResultListener oldListener = stream.getResultListener();
-                switch (params[0]) {
-                    case "file":
-                        if (!(oldListener instanceof DataWriter)) {
-                            Configuration config = Configuration.getConfig();
-
-                            String profilesFilename = config.getProperty(DataWriter.CFG_PROFILE_FILENAME);
-                            String tweetsFilename = config.getProperty(DataWriter.CFG_TWEETS_FILENAME);
-
-                            stream.setResultListener(new DataWriter(profilesFilename, tweetsFilename));
-
-                            // save the changes to the config.
-                            config.save();
-                        }
-                        break;
-                    case "shell":
-                        if (oldListener instanceof DataWriter) {
-                            ((DataWriter) oldListener).close();
-                        }
-                        if (!(oldListener instanceof StreamHandler)) {
-                            stream.setResultListener(new StreamHandler());
-                        }
-                        break;
-                    default:
-                        System.err.println("Unrecognized target " + params[0]);
-                        break;
+                CompositeResultListener rls;
+                rls = (CompositeResultListener) stream.getResultListener();
+                if (params.length > 0) {
+                    configureTargets(rls, params);
+                } else {
+                    ClassEnabledTracker<ResultListener> targets = getPossibleTargets(rls);
+                    // print the names of all targets that are enabled
+                    System.out.print("Enabled targets:");
+                    for (String name : targets.getNames()) {
+                        System.out.print(" " + name);
+                    }
+                    System.out.println();
                 }
                 break;
             default:
                 throw new AssertionError(command.name());
         }
+    }
+
+    /**
+     * @param stream An instance for which open targets apply to.
+     * @return All targets that can be disabled or enabled.
+     */
+    private ClassEnabledTracker<ResultListener> getPossibleTargets(CompositeResultListener rls) {
+        Map<String, Class<? extends ResultListener>> targets = new TreeMap<>();
+        targets.put("file", DataWriter.class);
+        targets.put("shell", StreamHandler.class);
+
+        ClassEnabledTracker<ResultListener> targetFoo = new ClassEnabledTracker<>(targets);
+        targetFoo.disableAll();
+        targetFoo.enableClasses(rls.getRegistered());
+        return targetFoo;
+    }
+
+    /**
+     * Process enable and disable target parameters.
+     */
+    private void configureTargets(CompositeResultListener rls, String[] params) {
+        ClassEnabledTracker<ResultListener> targets = getPossibleTargets(rls);
+
+        for (String type : params) {
+            String name;
+            // whether to remove, add targets or restrict to one
+            if (type.startsWith("-") || type.startsWith("+")) {
+                name = type.substring(1);
+            } else {
+                name = type;
+            }
+            if (!targets.has(name)) {
+                System.err.println("Unrecognized target: " + name);
+                continue;
+            }
+            // queue for enable or disable
+            if (type.startsWith("-")) {
+                targets.disable(name);
+            } else if (type.startsWith("+")) {
+                targets.enable(name);
+            } else {
+                targets.disableAll();
+                targets.enable(name);
+            }
+        }
+        for (String name : targets.getDisabled()) {
+            if (disableTarget(rls, targets.getClassByName(name))) {
+                System.err.println("Disabled " + name);
+            }
+        }
+        for (String name : targets.getEnabled()) {
+            if (enableTarget(rls, targets.getClassByName(name))) {
+                System.err.println("Enabled " + name);
+            }
+        }
+    }
+
+    private boolean enableTarget(CompositeResultListener rls,
+            Class<? extends ResultListener> rlCls) {
+        ResultListener oldListener = rls.findListener(rlCls);
+
+        // don't add it again if already activated
+        if (oldListener != null) {
+            return false;
+        }
+
+        if (rlCls == DataWriter.class) {
+            Configuration config = Configuration.getConfig();
+
+            String profilesFilename = config.getProperty(DataWriter.CFG_PROFILE_FILENAME);
+            String tweetsFilename = config.getProperty(DataWriter.CFG_TWEETS_FILENAME);
+
+            rls.register(new DataWriter(profilesFilename, tweetsFilename));
 
+            // save the changes to the config.
+            config.save();
+        } else if (rlCls == StreamHandler.class) {
+            rls.register(new StreamHandler());
+        }
+        return true;
+    }
+
+    private boolean disableTarget(CompositeResultListener rls,
+            Class<? extends ResultListener> rlCls) {
+        ResultListener oldListener = rls.findListener(rlCls);
+
+        // no need for action if not activated
+        if (oldListener == null) {
+            return false;
+        }
+
+        rls.unregister(oldListener);
+        // do we need to cleanup something?
+        if (oldListener instanceof DataWriter) {
+            ((DataWriter) oldListener).close();
+        }
+        return true;
     }
 }
diff --git a/src/provider/CompositeResultListener.java b/src/provider/CompositeResultListener.java
new file mode 100644
index 0000000..3db91da
--- /dev/null
+++ b/src/provider/CompositeResultListener.java
@@ -0,0 +1,56 @@
+package provider;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.TreeSet;
+import org.json.JSONObject;
+
+/**
+ * ResultListener that can contain multiple registrations.
+ *
+ * @author Peter Wu
+ */
+public class CompositeResultListener implements ResultListener {
+
+    private final List<ResultListener> listeners;
+
+    public CompositeResultListener() {
+        this.listeners = new ArrayList<>();
+    }
+
+    public void register(ResultListener rl) {
+        listeners.add(rl);
+    }
+
+    public void unregister(ResultListener rl) {
+        listeners.remove(rl);
+    }
+
+    public List<ResultListener> getRegistered() {
+        return new ArrayList<>(listeners);
+    }
+
+    public ResultListener findListener(Class<? extends ResultListener> rlCls) {
+        for (ResultListener rl : listeners) {
+            if (rlCls.isInstance(rl)) {
+                return rl;
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public void tweetGenerated(JSONObject obj) {
+        for (ResultListener rl : listeners) {
+            rl.tweetGenerated(obj);
+        }
+    }
+
+    @Override
+    public void profileGenerated(JSONObject obj) {
+        for (ResultListener rl : listeners) {
+            rl.profileGenerated(obj);
+        }
+    }
+}
diff --git a/src/support/ClassEnabledTracker.java b/src/support/ClassEnabledTracker.java
new file mode 100644
index 0000000..0859000
--- /dev/null
+++ b/src/support/ClassEnabledTracker.java
@@ -0,0 +1,95 @@
+package support;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.TreeSet;
+
+/**
+ * Tracks classes by name and whether those classes are enabled or not.
+ *
+ * @author Peter Wu
+ * @param <T> The type of classes that are to be tracked.
+ */
+public class ClassEnabledTracker<T> {
+
+    private final Map<String, Class<? extends T>> allowedClasses;
+    private final Map<String, Boolean> enabledClasses;
+
+    public ClassEnabledTracker(Map<String, Class<? extends T>> allowedClasses) {
+        this.allowedClasses = allowedClasses;
+        this.enabledClasses = new TreeMap<>();
+        for (String name : getNames()) {
+            enabledClasses.put(name, false);
+        }
+    }
+
+    public Set<String> getNames() {
+        return allowedClasses.keySet();
+    }
+
+    public Set<String> getEnabled() {
+        Set<String> enabled = new TreeSet<>();
+        for (Map.Entry<String, Boolean> e : enabledClasses.entrySet()) {
+            if (e.getValue()) {
+                enabled.add(e.getKey());
+            }
+        }
+        return enabled;
+    }
+
+    public Set<String> getDisabled() {
+        Set<String> disabled = new TreeSet<>(getNames());
+        disabled.removeAll(getEnabled());
+        return disabled;
+    }
+
+    public String getNameByClass(Class<T> aClass) {
+        for (Map.Entry<String, Class<? extends T>> e : allowedClasses.entrySet()) {
+            if (e.getValue() == aClass) {
+                return e.getKey();
+            }
+        }
+        return null;
+    }
+
+    public Class<? extends T> getClassByName(String name) {
+        return allowedClasses.get(name);
+    }
+
+    public boolean has(String name) {
+        return allowedClasses.containsKey(name);
+    }
+
+    public void disable(String name) {
+        enabledClasses.put(name, false);
+    }
+
+    public void enable(String name) {
+        enabledClasses.put(name, true);
+    }
+
+    public void disableAll() {
+        setEnabledAll(false);
+    }
+
+    public void enableAll() {
+        setEnabledAll(true);
+    }
+
+    private void setEnabledAll(boolean enabled) {
+        for (Map.Entry<String, Boolean> e : enabledClasses.entrySet()) {
+            e.setValue(enabled);
+        }
+    }
+
+    public void enableClasses(List<T> classes) {
+        for (T t : classes) {
+            String name = getNameByClass((Class<T>) t.getClass());
+            if (name != null) {
+                enabledClasses.put(name, true);
+            }
+        }
+    }
+}