From 0a06afab203f3992a37ced49500d9fadf1d299cd Mon Sep 17 00:00:00 2001
From: Steven Xie <dev@stevenxie.me>
Date: Sun, 12 Aug 2018 01:12:02 -0400
Subject: [PATCH] Now auto-adds trailing newline if necessary

A "-p" / "--preserve" flag has been added for times when one needs to
preserve the exact whitespacing of the source string.
---
 .gitignore |  3 ++-
 flags.go   |  3 ++-
 main.go    | 20 ++++++++++++++------
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index e827f72..bd624bc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
 ## ----- REPOSITORY -----
-## The compiled binary
+## The compiled binary (excluding Go source files)
 /dgen*
+!/dgen*.go
 
 ## ---- CODECOV -----
 coverage.txt
diff --git a/flags.go b/flags.go
index 2ec6c45..b510b8d 100644
--- a/flags.go
+++ b/flags.go
@@ -10,7 +10,8 @@ var (
 	// Opts are flag-enabled options for dgen.
 	Opts struct {
 		// Stats will show statistics at the end of the string dump.
-		Stats bool `short:"s" long:"stats" description:"Show statistics after string dump."`
+		Stats    bool `short:"s" long:"stats" description:"Show statistics after string dump."`
+		Preserve bool `short:"p" long:"preserve" description:"Preserve whitespacing; do not add terminating newlines."`
 	}
 
 	fparser = flags.NewParser(&Opts, flags.Default)
diff --git a/main.go b/main.go
index 9daf406..acac7ee 100644
--- a/main.go
+++ b/main.go
@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"github.com/steven-xie/dgen/throughput"
 	"os"
-	"strings"
 )
 
 func main() {
@@ -13,16 +12,17 @@ func main() {
 
 	n, err := throughput.Dump(str, reps, Bufsize, os.Stdout)
 
+	hasnl := hasTrailingNewline(str)
 	// Ensure that if extra info is about to be produced, there are at least two
 	// newlines before that info is printed.
 	if err != nil || Opts.Stats {
-		nlcount := strings.Count(str, "\n")
-		switch nlcount {
-		case 0:
-			fmt.Print("\n\n")
-		case 1:
+		if hasnl {
 			fmt.Print("\n")
+		} else {
+			fmt.Print("\n\n")
 		}
+	} else if !Opts.Preserve && !hasnl {
+		fmt.Print("\n")
 	}
 
 	if err != nil {
@@ -34,3 +34,11 @@ func main() {
 		fmt.Printf("Successfully printed %d bytes.\n", n)
 	}
 }
+
+func hasTrailingNewline(s string) bool {
+	strlen := len(s)
+	if strlen < 1 {
+		return false
+	}
+	return s[strlen-1] == '\n'
+}