Merge pull request #128 from Bobo1239/serve-squashed

Implement Serve feature

Author: azerupi

Cargo.toml

@@ -25,18 +25,24 @@ notify = { version = "2.5.5", optional = true }
 time = { version = "0.1.34", optional = true }
 crossbeam = { version = "0.2.8", optional = true }
 
+# Serve feature
+iron = { version = "0.3", optional = true }
+staticfile = { version = "0.2", optional = true }
+ws = { version = "0.4.6", optional = true}
+
 
 # Tests
 [dev-dependencies]
 tempdir = "0.3.4"
 
 
 [features]
-default = ["output", "watch"]
+default = ["output", "watch", "serve"]
 debug = []
 output = []
 regenerate-css = []
 watch = ["notify", "time", "crossbeam"]
+serve = ["iron", "staticfile", "ws"]
 
 [[bin]]
 doc = false

README.md

@@ -1,4 +1,4 @@
-# mdBook  
+# mdBook
 
 <table>
     <tr>
@@ -100,6 +100,10 @@ Here are the main commands you will want to run, for a more exhaustive explanati
 
     When you run this command, mdbook will watch your markdown files to rebuild the book on every change. This avoids having to come back to the terminal to type `mdbook build` over and over again.
 
+- `mdbook serve`
+
+    Does the same thing as `mdbook watch` but additionally serves the book at `http://localhost:3000` (port is changeable) and reloads the browser when a change occures.
+
 ### As a library
 
 Aside from the command line interface, this crate can also be used as a library. This means that you could integrate it in an existing project, like a web-app for example. Since the command line interface is just a wrapper around the library functionality, when you use this crate as a library you have full access to all the functionality of the command line interface with and easy to use API and more!

src/bin/mdbook.rs

@@ -10,6 +10,13 @@ extern crate notify;
 #[cfg(feature = "watch")]
 extern crate time;
 
+// Dependencies for the Serve feature
+#[cfg(feature = "serve")]
+extern crate iron;
+#[cfg(feature = "serve")]
+extern crate staticfile;
+#[cfg(feature = "serve")]
+extern crate ws;
 
 use std::env;
 use std::error::Error;
@@ -50,6 +57,11 @@ fn main() {
                     .subcommand(SubCommand::with_name("watch")
                         .about("Watch the files for changes")
                         .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'"))
+                    .subcommand(SubCommand::with_name("serve")
+                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
+                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'")
+                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
+                        .arg_from_usage("-w, --websocket-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'"))
                     .subcommand(SubCommand::with_name("test")
                         .about("Test that code samples compile"))
                     .get_matches();
@@ -60,6 +72,8 @@ fn main() {
         ("build", Some(sub_matches)) => build(sub_matches),
         #[cfg(feature = "watch")]
         ("watch", Some(sub_matches)) => watch(sub_matches),
+        #[cfg(feature = "serve")]
+        ("serve", Some(sub_matches)) => serve(sub_matches),
         ("test", Some(sub_matches)) => test(sub_matches),
         (_, _) => unreachable!(),
     };
@@ -148,77 +162,85 @@ fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
 #[cfg(feature = "watch")]
 fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
     let book_dir = get_book_dir(args);
-    let book = MDBook::new(&book_dir).read_config();
+    let mut book = MDBook::new(&book_dir).read_config();
 
-    // Create a channel to receive the events.
-    let (tx, rx) = channel();
+    trigger_on_change(&mut book, |event, book| {
+        if let Some(path) = event.path {
+            println!("File changed: {:?}\nBuilding book...\n", path);
+            match book.build() {
+                Err(e) => println!("Error while building: {:?}", e),
+                _ => {},
+            }
+            println!("");
+        }
+    });
 
-    let w: Result<notify::RecommendedWatcher, notify::Error> = notify::Watcher::new(tx);
+    Ok(())
+}
 
-    match w {
-        Ok(mut watcher) => {
 
-            // Add the source directory to the watcher
-            if let Err(e) = watcher.watch(book.get_src()) {
-                println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
-                ::std::process::exit(0);
-            };
+// Watch command implementation
+#[cfg(feature = "serve")]
+fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
+    const RELOAD_COMMAND: &'static str = "reload";
 
-            // Add the book.json file to the watcher if it exists, because it's not
-            // located in the source directory
-            if let Err(_) = watcher.watch(book_dir.join("book.json")) {
-                // do nothing if book.json is not found
-            }
+    let book_dir = get_book_dir(args);
+    let mut book = MDBook::new(&book_dir).read_config();
+    let port = args.value_of("port").unwrap_or("3000");
+    let ws_port = args.value_of("ws-port").unwrap_or("3001");
+
+    let address = format!("localhost:{}", port);
+    let ws_address = format!("localhost:{}", ws_port);
+
+    book.set_livereload(format!(r#"
+        <script type="text/javascript">
+            var socket = new WebSocket("ws://localhost:{}");
+            socket.onmessage = function (event) {{
+                if (event.data === "{}") {{
+                    socket.close();
+                    location.reload(true); // force reload from server (not from cache)
+                }}
+            }};
+
+            window.onbeforeunload = function() {{
+                socket.close();
+            }}
+        </script>
+    "#, ws_port, RELOAD_COMMAND).to_owned());
 
-            let mut previous_time = time::get_time();
+    try!(book.build());
 
-            crossbeam::scope(|scope| {
-                loop {
-                    match rx.recv() {
-                        Ok(event) => {
-
-                            // Skip the event if an event has already been issued in the last second
-                            let time = time::get_time();
-                            if time - previous_time < time::Duration::seconds(1) {
-                                continue;
-                            } else {
-                                previous_time = time;
-                            }
-
-                            if let Some(path) = event.path {
-                                // Trigger the build process in a new thread (to keep receiving events)
-                                scope.spawn(move || {
-                                    println!("File changed: {:?}\nBuilding book...\n", path);
-                                    match build(args) {
-                                        Err(e) => println!("Error while building: {:?}", e),
-                                        _ => {},
-                                    }
-                                    println!("");
-                                });
-
-                            } else {
-                                continue;
-                            }
-                        },
-                        Err(e) => {
-                            println!("An error occured: {:?}", e);
-                        },
-                    }
-                }
-            });
+    let staticfile = staticfile::Static::new(book.get_dest());
+    let iron = iron::Iron::new(staticfile);
+    let _iron = iron.http(&*address).unwrap();
 
-        },
-        Err(e) => {
-            println!("Error while trying to watch the files:\n\n\t{:?}", e);
-            ::std::process::exit(0);
-        },
-    }
+    let ws_server = ws::WebSocket::new(|_| {
+        |_| {
+            Ok(())
+        }
+    }).unwrap();
+
+    let broadcaster = ws_server.broadcaster();
+
+    std::thread::spawn(move || {
+        ws_server.listen(&*ws_address).unwrap();
+    });
+
+    trigger_on_change(&mut book, move |event, book| {
+        if let Some(path) = event.path {
+            println!("File changed: {:?}\nBuilding book...\n", path);
+            match book.build() {
+                Err(e) => println!("Error while building: {:?}", e),
+                _ => broadcaster.send(RELOAD_COMMAND).unwrap(),
+            }
+            println!("");
+        }
+    });
 
     Ok(())
 }
 
 
-
 fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
     let book_dir = get_book_dir(args);
     let mut book = MDBook::new(&book_dir).read_config();
@@ -229,7 +251,6 @@ fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
 }
 
 
-
 fn get_book_dir(args: &ArgMatches) -> PathBuf {
     if let Some(dir) = args.value_of("dir") {
         // Check if path is relative from current dir, or absolute...
@@ -243,3 +264,57 @@ fn get_book_dir(args: &ArgMatches) -> PathBuf {
         env::current_dir().unwrap()
     }
 }
+
+
+// Calls the closure when a book source file is changed. This is blocking!
+fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
+    where F: Fn(notify::Event, &mut MDBook) -> ()
+{
+    // Create a channel to receive the events.
+    let (tx, rx) = channel();
+
+    let w: Result<notify::RecommendedWatcher, notify::Error> = notify::Watcher::new(tx);
+
+    match w {
+        Ok(mut watcher) => {
+            // Add the source directory to the watcher
+            if let Err(e) = watcher.watch(book.get_src()) {
+                println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
+                ::std::process::exit(0);
+            };
+
+            // Add the book.json file to the watcher if it exists, because it's not
+            // located in the source directory
+            if let Err(_) = watcher.watch(book.get_root().join("book.json")) {
+                // do nothing if book.json is not found
+            }
+
+            let mut previous_time = time::get_time();
+
+            println!("\nListening for changes...\n");
+
+            loop {
+                match rx.recv() {
+                    Ok(event) => {
+                        // Skip the event if an event has already been issued in the last second
+                        let time = time::get_time();
+                        if time - previous_time < time::Duration::seconds(1) {
+                            continue;
+                        } else {
+                            previous_time = time;
+                        }
+
+                        closure(event, book);
+                    },
+                    Err(e) => {
+                        println!("An error occured: {:?}", e);
+                    },
+                }
+            }
+        },
+        Err(e) => {
+            println!("Error while trying to watch the files:\n\n\t{:?}", e);
+            ::std::process::exit(0);
+        },
+    }
+}

src/book/mdbook.rs

@@ -15,6 +15,8 @@ pub struct MDBook {
     config: BookConfig,
     pub content: Vec<BookItem>,
     renderer: Box<Renderer>,
+    #[cfg(feature = "serve")]
+    livereload: Option<String>,
 }
 
 impl MDBook {
@@ -38,6 +40,7 @@ impl MDBook {
                         .set_dest(&root.join("book"))
                         .to_owned(),
             renderer: Box::new(HtmlHandlebars::new()),
+            livereload: None,
         }
     }
 
@@ -398,6 +401,23 @@ impl MDBook {
         &self.config.description
     }
 
+    pub fn set_livereload(&mut self, livereload: String) -> &mut Self {
+        self.livereload = Some(livereload);
+        self
+    }
+
+    pub fn unset_livereload(&mut self) -> &Self {
+        self.livereload = None;
+        self
+    }
+
+    pub fn get_livereload(&self) -> Option<&String> {
+        match self.livereload {
+            Some(ref livereload) => Some(&livereload),
+            None => None,
+        }
+    }
+
     // Construct book
     fn parse_summary(&mut self) -> Result<(), Box<Error>> {
         // When append becomes stable, use self.content.append() ...

src/renderer/html_handlebars/hbs_renderer.rs

@@ -287,6 +287,9 @@ fn make_data(book: &MDBook) -> Result<BTreeMap<String, Json>, Box<Error>> {
     data.insert("title".to_owned(), book.get_title().to_json());
     data.insert("description".to_owned(), book.get_description().to_json());
     data.insert("favicon".to_owned(), "favicon.png".to_json());
+    if let Some(livereload) = book.get_livereload() {
+        data.insert("livereload".to_owned(), livereload.to_json());
+    }
 
     let mut chapters = vec![];
 

src/theme/index.hbs

@@ -107,6 +107,9 @@
             }
         </script>
 
+        <!-- Livereload script (if served using the cli tool) -->
+        {{{livereload}}}
+
         <script src="highlight.js"></script>
         <script src="book.js"></script>
     </body>

src/utils/fs.rs

@@ -150,7 +150,7 @@ pub fn copy_files_except_ext(from: &Path, to: &Path, recursive: bool, ext_blackl
                 debug!("[*] creating path for file: {:?}",
                        &to.join(entry.path().file_name().expect("a file should have a file name...")));
 
-                output!("[*] copying file: {:?}\n    to {:?}",
+                output!("[*] Copying file: {:?}\n    to {:?}",
                         entry.path(),
                         &to.join(entry.path().file_name().expect("a file should have a file name...")));
                 try!(fs::copy(entry.path(),