Merge pull request #5 from marcesher/add_depth

Add support for a --depth flag

README.md

@@ -63,6 +63,7 @@ usage: gh-md-toc [<flags>] [<path>...]
 Flags:
   --help     Show help (also see --help-long and --help-man).
   --version  Show application version.
+  --depth    How many levels of headings to include. Defaults to 0 (all)
 
 Args:
   [<path>]  Local path or URL of the document to grab TOC
@@ -114,6 +115,7 @@ Table of Contents
         * [OR using Vundle:](#or-using-vundle)
   * [License](#license)
 ```
+
 Remote files
 ------------
 
@@ -252,6 +254,23 @@ You can easily combine both ways:
 Created by [gh-md-toc](https://github.com/ekalinin/github-markdown-toc)
 ```
 
+Depth
+-----
+
+Use `--depth=INT` to control how many levels of headers to include in the TOC
+
+```bash
+➥ ./gh-md-toc --depth=1 ~/projects/Dockerfile.vim/README.md                                                                                                                                                Вс. марта 22 22:51:46 MSK 2015
+
+Table of Contents
+=================
+
+  * [Dockerfile.vim](#dockerfilevim)
+  * [Screenshot](#screenshot)
+  * [Installation](#installation)
+  * [License](#license)
+```
+
 
 LICENSE
 =======

main.go

@@ -22,6 +22,9 @@ var (
 )
 
 type GHToc []string
+type Options struct {
+	Depth int
+}
 
 //
 // Internal
@@ -130,11 +133,11 @@ func ConvertMd2Html(localpath string) string {
 }
 
 // Create TOC by html from github
-func GrabToc(html string) *GHToc {
-	return GrabTocX(html, "")
+func GrabToc(html string, opts Options) *GHToc {
+	return GrabTocX(html, "", opts)
 }
 
-func GrabTocX(html string, absPath string) *GHToc {
+func GrabTocX(html string, absPath string, opts Options) *GHToc {
 	re := `(?si)<h(?P<num>[1-6])>\s*` +
 		`<a\s*id="user-content-[^"]*"\s*class="anchor"\s*` +
 		`href="(?P<href>[^"]*)"[^>]*>\s*` +
@@ -153,6 +156,10 @@ func GrabTocX(html string, absPath string) *GHToc {
 		}
 		// format result
 		n, _ := strconv.Atoi(groups["num"])
+		if opts.Depth > 0 && n > opts.Depth {
+			continue
+		}
+
 		link := groups["href"]
 		if len(absPath) > 0 {
 			link = absPath + link
@@ -168,16 +175,16 @@ func GrabTocX(html string, absPath string) *GHToc {
 }
 
 // Generate TOC for document (path in filesystem or url)
-func GenerateToc(path string) *GHToc {
-	return GenerateTocX(path, false)
+func GenerateToc(path string, opts Options) *GHToc {
+	return GenerateTocX(path, false, opts)
 }
 
-func GenerateTocX(path string, absPaths bool) *GHToc {
+func GenerateTocX(path string, absPaths bool, opts Options) *GHToc {
 	htmlBody := GetHmtlBody(path)
 	if absPaths {
-		return GrabTocX(htmlBody, path)
+		return GrabTocX(htmlBody, path, opts)
 	} else {
-		return GrabToc(htmlBody)
+		return GrabToc(htmlBody, opts)
 	}
 }
 
@@ -195,9 +202,14 @@ func main() {
 		"If not entered, then read Markdown from stdin."
 	paths := kingpin.Arg("path", paths_desc).Strings()
 	serial := kingpin.Flag("serial", "Grab TOCs in the serial mode").Bool()
+	depth := kingpin.Flag("depth", "How many levels of headings to include. Defaults to 0 (all)").Default("0").Int()
 	kingpin.Version(version)
 	kingpin.Parse()
 
+	opts := Options{
+		Depth: *depth,
+	}
+
 	pathsCount := len(*paths)
 
 	if pathsCount == 1 {
@@ -212,10 +224,10 @@ func main() {
 	ch := make(chan *GHToc, pathsCount)
 	for _, p := range *paths {
 		if *serial {
-			ch <- GenerateTocX(p, absPathsInToc)
+			ch <- GenerateTocX(p, absPathsInToc, opts)
 		} else {
 			go func(path string, showAbsPath bool) {
-				ch <- GenerateTocX(path, absPathsInToc)
+				ch <- GenerateTocX(path, absPathsInToc, opts)
 			}(p, absPathsInToc)
 		}
 	}
@@ -234,7 +246,7 @@ func main() {
 		bytes, err := ioutil.ReadAll(os.Stdin)
 		check(err)
 		check(ioutil.WriteFile(file.Name(), bytes, 0644))
-		PrintToc(GenerateToc(file_path))
+		PrintToc(GenerateToc(file_path, opts))
 	}
 
 	fmt.Println("Created by [gh-md-toc](https://github.com/ekalinin/github-markdown-toc.go)")

main_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"testing"
 )
 
@@ -22,7 +23,7 @@ func Test_grab_toc_onerow(t *testing.T) {
 	}
 	toc := *GrabToc(`
 	<h1><a id="user-content-readme-in-another-language" class="anchor" href="#readme-in-another-language" aria-hidden="true"><span class="octicon octicon-link"></span></a>README in another language</h1>
-	`)
+	`, Options{})
 	if toc[0] != toc_expected[0] {
 		t.Error("Res :", toc, "\nExpected     :", toc_expected)
 	}
@@ -39,7 +40,7 @@ func Test_grab_toc_onerow_with_newlines(t *testing.T) {
 		</a>
 		README in another language
 	</h1>
-	`)
+	`, Options{})
 	if toc[0] != toc_expected[0] {
 		t.Error("Res :", toc, "\nExpected     :", toc_expected)
 	}
@@ -69,7 +70,7 @@ case you need to implement 2 functions in the plugin's body:</p>
 
 <p>This function should return list of available versions of the plugin.
 For example:</p>
-	`)
+	`, Options{})
 	for i := 0; i <= len(toc_expected)-1; i++ {
 		if toc[i] != toc_expected[i] {
 			t.Error("Res :", toc[i], "\nExpected     :", toc_expected[i])
@@ -107,7 +108,46 @@ func Test_GrabToc_backquoted(t *testing.T) {
 <a id="user-content-the-command-bar2-is-better" class="anchor" href="#the-command-bar2-is-better" aria-hidden="true"><span class="octicon octicon-link"></span></a>The command <code>bar2</code> is better</h2>
 
 <p>Blabla...</p>
-	`)
+	`, Options{})
+
+	for i := 0; i <= len(toc_expected)-1; i++ {
+		if toc[i] != toc_expected[i] {
+			t.Error("Res :", toc[i], "\nExpected      :", toc_expected[i])
+		}
+	}
+}
+
+func Test_GrabToc_depth(t *testing.T) {
+	toc_expected := []string{
+		"  * [The command foo1](#the-command-foo1)",
+		"  * [The command bar1](#the-command-bar1)",
+	}
+
+	toc := *GrabToc(`
+<h1>
+<a id="user-content-the-command-foo1" class="anchor" href="#the-command-foo1" aria-hidden="true"><span class="octicon octicon-link"></span></a>The command <code>foo1</code>
+</h1>
+
+<p>Blabla...</p>
+
+<h2>
+<a id="user-content-the-command-foo2-is-better" class="anchor" href="#the-command-foo2-is-better" aria-hidden="true"><span class="octicon octicon-link"></span></a>The command <code>foo2</code> is better</h2>
+
+<p>Blabla...</p>
+
+<h1>
+<a id="user-content-the-command-bar1" class="anchor" href="#the-command-bar1" aria-hidden="true"><span class="octicon octicon-link"></span></a>The command <code>bar1</code>
+</h1>
+
+<p>Blabla...</p>
+
+<h2>
+<a id="user-content-the-command-bar2-is-better" class="anchor" href="#the-command-bar2-is-better" aria-hidden="true"><span class="octicon octicon-link"></span></a>The command <code>bar2</code> is better</h2>
+
+<p>Blabla...</p>
+	`, Options{Depth: 1})
+
+	fmt.Println(toc)
 
 	for i := 0; i <= len(toc_expected)-1; i++ {
 		if toc[i] != toc_expected[i] {
@@ -123,7 +163,7 @@ func Test_grab_toc_with_abspath(t *testing.T) {
 	}
 	toc := *GrabTocX(`
 	<h1><a id="user-content-readme-in-another-language" class="anchor" href="#readme-in-another-language" aria-hidden="true"><span class="octicon octicon-link"></span></a>README in another language</h1>
-	`, link)
+	`, link, Options{})
 	if toc[0] != toc_expected[0] {
 		t.Error("Res :", toc, "\nExpected     :", toc_expected)
 	}
@@ -141,7 +181,7 @@ func Test_EscapedChars(t *testing.T) {
 				<span class="octicon octicon-link"></span>
 			</a>
 			mod_*
-		</h2>`)
+		</h2>`, Options{})
 
 	if toc[0] != toc_expected[0] {
 		t.Error("Res :", toc, "\nExpected     :", toc_expected)