Update

README.md

@@ -29,6 +29,15 @@ Markdown Table Of Contents is a straightforward command-line tool designed to ge
 
 3. **Output**: The tool will generate a table of contents based on the headings in your Markdown file.
 
+**Arguments:**
+
+```txt
+  -depth int
+        Depth of the table of contents (1,2,...,6). 2 will print only H1 and H2 (default 3)
+  -indent string
+        Indentation string ('  ','tab') (default "  ")
+```
+
 ## Example
 
 Suppose you have a Markdown file named `example.md` with the following headings:
@@ -63,6 +72,4 @@ This project is licensed under the [MIT License](LICENSE).
 
 ## Todolist
 
-- If no H1, or H2, ..., trim spaces before
-- Parse command line arguments
 - Replace from file with {table_of_contents} tag or something

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"os"
 	"regexp"
@@ -18,33 +19,64 @@ func formatTitle(title string) string {
 }
 
 func main() {
-	if len(os.Args) <= 1 {
-		fmt.Println("Usage: markdown-table-of-contents [markdownfile.md] <Identation string>")
-		os.Exit(1)
-	}
+	var indentationString string
+	var depth int
+
+	flag.StringVar(&indentationString, "indent", "  ", "Indentation string ('  ','tab')")
+	flag.IntVar(&depth, "depth", 3, "Depth of the table of contents (1,2,...,6). 2 will print only H1 and H2")
+
+	flag.Parse()
 
-	indentationString := "  " // By default, use 2 spaces
-	if len(os.Args) > 2 {
-		indentationString = os.Args[2]
+	var filePath string
+	if len(flag.Args()) < 1 {
+		fmt.Println("Usage: markdown-table-of-contents [markdownfile.md]")
+		fmt.Println("Type 'markdown-table-of-contents -h' for help")
+		os.Exit(1)
 	}
+	filePath = flag.Args()[0]
 
-	filePath := os.Args[1]
 	md := markdown.New(filePath)
 	err := md.Read()
 	if err != nil {
 		fmt.Println(err)
 		os.Exit(1)
 	}
-	str := ""
+
+	var strs []string = []string{}
 	for _, section := range md.Sections {
 		if section.SectionType == markdown.NullSection {
 			continue
 		}
 
 		numberOfSpaces := len(string(section.SectionType)) - 1
+		if numberOfSpaces+1 > depth {
+			continue
+		}
+		var str string
 		str += strings.Repeat(indentationString, numberOfSpaces)
 		str += "- [" + section.Text + "](#" + formatTitle(section.Text) + ")"
-		str += "\n"
+		strs = append(strs, str)
+	}
+
+	// If all start with indent, remove all
+	for {
+		allStartWithIndent := true
+		for _, str := range strs {
+			if !strings.HasPrefix(str, indentationString) {
+				allStartWithIndent = false
+				break
+			}
+		}
+		if allStartWithIndent {
+			for i := range strs {
+				strs[i] = strs[i][len(indentationString):]
+			}
+		} else {
+			break
+		}
+	}
+
+	for _, str := range strs {
+		fmt.Println(str)
 	}
-	fmt.Println(str)
 }