Merge pull request #803 from TheHipbot/branch-tracking-on-clone

config: adds branches to config for tracking branches against remotes…

config/branch.go

···
+package config
+
+import (
+	"errors"
+
+	"gopkg.in/src-d/go-git.v4/plumbing"
+	format "gopkg.in/src-d/go-git.v4/plumbing/format/config"
+)
+
+var (
+	errBranchEmptyName    = errors.New("branch config: empty name")
+	errBranchInvalidMerge = errors.New("branch config: invalid merge")
+)
+
+// Branch contains information on the
+// local branches and which remote to track
+type Branch struct {
+	// Name of branch
+	Name string
+	// Remote name of remote to track
+	Remote string
+	// Merge is the local refspec for the branch
+	Merge plumbing.ReferenceName
+
+	raw *format.Subsection
+}
+
+// Validate validates fields of branch
+func (b *Branch) Validate() error {
+	if b.Name == "" {
+		return errBranchEmptyName
+	}
+
+	if b.Merge != "" && !b.Merge.IsBranch() {
+		return errBranchInvalidMerge
+	}
+
+	return nil
+}
+
+func (b *Branch) marshal() *format.Subsection {
+	if b.raw == nil {
+		b.raw = &format.Subsection{}
+	}
+
+	b.raw.Name = b.Name
+
+	if b.Remote == "" {
+		b.raw.RemoveOption(remoteSection)
+	} else {
+		b.raw.SetOption(remoteSection, b.Remote)
+	}
+
+	if b.Merge == "" {
+		b.raw.RemoveOption(mergeKey)
+	} else {
+		b.raw.SetOption(mergeKey, string(b.Merge))
+	}
+
+	return b.raw
+}
+
+func (b *Branch) unmarshal(s *format.Subsection) error {
+	b.raw = s
+
+	b.Name = b.raw.Name
+	b.Remote = b.raw.Options.Get(remoteSection)
+	b.Merge = plumbing.ReferenceName(b.raw.Options.Get(mergeKey))
+
+	return b.Validate()
+}

config/branch_test.go

···
+package config
+
+import (
+	. "gopkg.in/check.v1"
+	"gopkg.in/src-d/go-git.v4/plumbing"
+)
+
+type BranchSuite struct{}
+
+var _ = Suite(&BranchSuite{})
+
+func (b *BranchSuite) TestValidateName(c *C) {
+	goodBranch := Branch{
+		Name:   "master",
+		Remote: "some_remote",
+		Merge:  "refs/heads/master",
+	}
+	badBranch := Branch{
+		Remote: "some_remote",
+		Merge:  "refs/heads/master",
+	}
+	c.Assert(goodBranch.Validate(), IsNil)
+	c.Assert(badBranch.Validate(), NotNil)
+}
+
+func (b *BranchSuite) TestValidateMerge(c *C) {
+	goodBranch := Branch{
+		Name:   "master",
+		Remote: "some_remote",
+		Merge:  "refs/heads/master",
+	}
+	badBranch := Branch{
+		Name:   "master",
+		Remote: "some_remote",
+		Merge:  "blah",
+	}
+	c.Assert(goodBranch.Validate(), IsNil)
+	c.Assert(badBranch.Validate(), NotNil)
+}
+
+func (b *BranchSuite) TestMarshall(c *C) {
+	expected := []byte(`[core]
+	bare = false
+[branch "branch-tracking-on-clone"]
+	remote = fork
+	merge = refs/heads/branch-tracking-on-clone
+`)
+
+	cfg := NewConfig()
+	cfg.Branches["branch-tracking-on-clone"] = &Branch{
+		Name:   "branch-tracking-on-clone",
+		Remote: "fork",
+		Merge:  plumbing.ReferenceName("refs/heads/branch-tracking-on-clone"),
+	}
+
+	actual, err := cfg.Marshal()
+	c.Assert(err, IsNil)
+	c.Assert(string(actual), Equals, string(expected))
+}
+
+func (b *BranchSuite) TestUnmarshall(c *C) {
+	input := []byte(`[core]
+	bare = false
+[branch "branch-tracking-on-clone"]
+	remote = fork
+	merge = refs/heads/branch-tracking-on-clone
+`)
+
+	cfg := NewConfig()
+	err := cfg.Unmarshal(input)
+	c.Assert(err, IsNil)
+	branch := cfg.Branches["branch-tracking-on-clone"]
+	c.Assert(branch.Name, Equals, "branch-tracking-on-clone")
+	c.Assert(branch.Remote, Equals, "fork")
+	c.Assert(branch.Merge, Equals, plumbing.ReferenceName("refs/heads/branch-tracking-on-clone"))
+}

config/config.go

···
 }
 
 var (
-	ErrInvalid               = errors.New("config invalid remote")
+	ErrInvalid               = errors.New("config invalid key in remote or branch")
 	ErrRemoteConfigNotFound  = errors.New("remote config not found")
 	ErrRemoteConfigEmptyURL  = errors.New("remote config: empty URL")
 	ErrRemoteConfigEmptyName = errors.New("remote config: empty name")
···
 	// Submodules list of repository submodules, the key of the map is the name
 	// of the submodule, should equal to Submodule.Name.
 	Submodules map[string]*Submodule
-
+	// Branches list of branches, the key is the branch name and should
+	// equal Branch.Name
+	Branches map[string]*Branch
 	// Raw contains the raw information of a config file. The main goal is
 	// preserve the parsed information from the original format, to avoid
 	// dropping unsupported fields.
···
 	config := &Config{
 		Remotes:    make(map[string]*RemoteConfig),
 		Submodules: make(map[string]*Submodule),
+		Branches:   make(map[string]*Branch),
 		Raw:        format.New(),
 	}
 
···
 		}
 
 		if err := r.Validate(); err != nil {
+			return err
+		}
+	}
+
+	for name, b := range c.Branches {
+		if b.Name != name {
+			return ErrInvalid
+		}
+
+		if err := b.Validate(); err != nil {
 			return err
 		}
 	}
···
 const (
 	remoteSection    = "remote"
 	submoduleSection = "submodule"
+	branchSection    = "branch"
 	coreSection      = "core"
 	packSection      = "pack"
 	fetchKey         = "fetch"
···
 	bareKey          = "bare"
 	worktreeKey      = "worktree"
 	windowKey        = "window"
+	mergeKey         = "merge"
 
 	// DefaultPackWindow holds the number of previous objects used to
 	// generate deltas. The value 10 is the same used by git command.
···
 		return err
 	}
 	c.unmarshalSubmodules()
+
+	if err := c.unmarshalBranches(); err != nil {
+		return err
+	}
+
 	return c.unmarshalRemotes()
 }
 
···
 	}
 }
 
+func (c *Config) unmarshalBranches() error {
+	bs := c.Raw.Section(branchSection)
+	for _, sub := range bs.Subsections {
+		b := &Branch{}
+
+		if err := b.unmarshal(sub); err != nil {
+			return err
+		}
+
+		c.Branches[b.Name] = b
+	}
+	return nil
+}
+
 // Marshal returns Config encoded as a git-config file.
 func (c *Config) Marshal() ([]byte, error) {
 	c.marshalCore()
 	c.marshalPack()
 	c.marshalRemotes()
 	c.marshalSubmodules()
+	c.marshalBranches()
 
 	buf := bytes.NewBuffer(nil)
 	if err := format.NewEncoder(buf).Encode(c.Raw); err != nil {
···
 		s.Subsections[i] = section
 		i++
 	}
+}
+
+func (c *Config) marshalBranches() {
+	s := c.Raw.Section(branchSection)
+	newSubsections := make(format.Subsections, 0, len(c.Branches))
+	added := make(map[string]bool)
+	for _, subsection := range s.Subsections {
+		if branch, ok := c.Branches[subsection.Name]; ok {
+			newSubsections = append(newSubsections, branch.marshal())
+			added[subsection.Name] = true
+		}
+	}
+
+	branchNames := make([]string, 0, len(c.Branches))
+	for name := range c.Branches {
+		branchNames = append(branchNames, name)
+	}
+
+	sort.Strings(branchNames)
+
+	for _, name := range branchNames {
+		if !added[name] {
+			newSubsections = append(newSubsections, c.Branches[name].marshal())
+		}
+	}
+
+	s.Subsections = newSubsections
 }
 
 // RemoteConfig contains the configuration for a given remote repository.

config/config_test.go

···
 package config
 
-import . "gopkg.in/check.v1"
+import (
+	. "gopkg.in/check.v1"
+	"gopkg.in/src-d/go-git.v4/plumbing"
+)
 
 type ConfigSuite struct{}
 
···
 	c.Assert(cfg.Submodules["qux"].Name, Equals, "qux")
 	c.Assert(cfg.Submodules["qux"].URL, Equals, "https://github.com/foo/qux.git")
 	c.Assert(cfg.Submodules["qux"].Branch, Equals, "bar")
-
+	c.Assert(cfg.Branches["master"].Remote, Equals, "origin")
+	c.Assert(cfg.Branches["master"].Merge, Equals, plumbing.ReferenceName("refs/heads/master"))
 }
 
 func (s *ConfigSuite) TestMarshall(c *C) {
···
 	url = [e-mail]:mcuadros/go-git.git
 [submodule "qux"]
 	url = https://github.com/foo/qux.git
+[branch "master"]
+	remote = origin
+	merge = refs/heads/master
 `)
 
 	cfg := NewConfig()
···
 	cfg.Submodules["qux"] = &Submodule{
 		Name: "qux",
 		URL:  "https://github.com/foo/qux.git",
+	}
+
+	cfg.Branches["master"] = &Branch{
+		Name:   "master",
+		Remote: "origin",
+		Merge:  "refs/heads/master",
 	}
 
 	b, err := cfg.Marshal()
···
 	c.Assert(string(output), DeepEquals, string(input))
 }
 
+func (s *ConfigSuite) TestValidateConfig(c *C) {
+	config := &Config{
+		Remotes: map[string]*RemoteConfig{
+			"bar": {
+				Name: "bar",
+				URLs: []string{"http://foo/bar"},
+			},
+		},
+		Branches: map[string]*Branch{
+			"bar": {
+				Name: "bar",
+			},
+			"foo": {
+				Name:   "foo",
+				Remote: "origin",
+				Merge:  plumbing.ReferenceName("refs/heads/foo"),
+			},
+		},
+	}
+
+	c.Assert(config.Validate(), IsNil)
+}
+
 func (s *ConfigSuite) TestValidateInvalidRemote(c *C) {
 	config := &Config{
 		Remotes: map[string]*RemoteConfig{
···
 	c.Assert(config.Validate(), Equals, ErrRemoteConfigEmptyURL)
 }
 
-func (s *ConfigSuite) TestValidateInvalidKey(c *C) {
+func (s *ConfigSuite) TestValidateInvalidRemoteKey(c *C) {
 	config := &Config{
 		Remotes: map[string]*RemoteConfig{
 			"bar": {Name: "foo"},
···
 	c.Assert(fetch[0].String(), Equals, "+refs/heads/*:refs/remotes/foo/*")
 }
 
+func (s *ConfigSuite) TestValidateInvalidBranchKey(c *C) {
+	config := &Config{
+		Branches: map[string]*Branch{
+			"foo": {
+				Name:   "bar",
+				Remote: "origin",
+				Merge:  plumbing.ReferenceName("refs/heads/bar"),
+			},
+		},
+	}
+
+	c.Assert(config.Validate(), Equals, ErrInvalid)
+}
+
+func (s *ConfigSuite) TestValidateInvalidBranch(c *C) {
+	config := &Config{
+		Branches: map[string]*Branch{
+			"bar": {
+				Name:   "bar",
+				Remote: "origin",
+				Merge:  plumbing.ReferenceName("refs/heads/bar"),
+			},
+			"foo": {
+				Name:   "foo",
+				Remote: "origin",
+				Merge:  plumbing.ReferenceName("baz"),
+			},
+		},
+	}
+
+	c.Assert(config.Validate(), Equals, errBranchInvalidMerge)
+}
+
 func (s *ConfigSuite) TestRemoteConfigDefaultValues(c *C) {
 	config := NewConfig()
 
 	c.Assert(config.Remotes, HasLen, 0)
+	c.Assert(config.Branches, HasLen, 0)
 	c.Assert(config.Submodules, HasLen, 0)
 	c.Assert(config.Raw, NotNil)
 	c.Assert(config.Pack.Window, Equals, DefaultPackWindow)

repository.go

···
 )
 
 var (
+	// ErrBranchExists an error stating the specified branch already exists
+	ErrBranchExists = errors.New("branch already exists")
+	// ErrBranchNotFound an error stating the specified branch does not exist
+	ErrBranchNotFound            = errors.New("branch not found")
 	ErrInvalidReference          = errors.New("invalid reference, should be a tag or a branch")
 	ErrRepositoryNotExists       = errors.New("repository does not exist")
 	ErrRepositoryAlreadyExists   = errors.New("repository already exists")
 	ErrRemoteNotFound            = errors.New("remote not found")
-	ErrRemoteExists              = errors.New("remote already exists	")
+	ErrRemoteExists              = errors.New("remote already exists")
 	ErrWorktreeNotProvided       = errors.New("worktree should be provided")
 	ErrIsBareRepository          = errors.New("worktree not available in a bare repository")
 	ErrUnableToResolveCommit     = errors.New("unable to resolve commit")
···
 	return r.Storer.SetConfig(cfg)
 }
 
+// Branch return a Branch if exists
+func (r *Repository) Branch(name string) (*config.Branch, error) {
+	cfg, err := r.Storer.Config()
+	if err != nil {
+		return nil, err
+	}
+
+	b, ok := cfg.Branches[name]
+	if !ok {
+		return nil, ErrBranchNotFound
+	}
+
+	return b, nil
+}
+
+// CreateBranch creates a new Branch
+func (r *Repository) CreateBranch(c *config.Branch) error {
+	if err := c.Validate(); err != nil {
+		return err
+	}
+
+	cfg, err := r.Storer.Config()
+	if err != nil {
+		return err
+	}
+
+	if _, ok := cfg.Branches[c.Name]; ok {
+		return ErrBranchExists
+	}
+
+	cfg.Branches[c.Name] = c
+	return r.Storer.SetConfig(cfg)
+}
+
+// DeleteBranch delete a Branch from the repository and delete the config
+func (r *Repository) DeleteBranch(name string) error {
+	cfg, err := r.Storer.Config()
+	if err != nil {
+		return err
+	}
+
+	if _, ok := cfg.Branches[name]; !ok {
+		return ErrBranchNotFound
+	}
+
+	delete(cfg.Branches, name)
+	return r.Storer.SetConfig(cfg)
+}
+
 func (r *Repository) resolveToCommitHash(h plumbing.Hash) (plumbing.Hash, error) {
 	obj, err := r.Storer.EncodedObject(plumbing.AnyObject, h)
 	if err != nil {
···
 		}
 	}
 
-	return r.updateRemoteConfigIfNeeded(o, c, ref)
+	if err := r.updateRemoteConfigIfNeeded(o, c, ref); err != nil {
+		return err
+	}
+
+	if ref.Name().IsBranch() {
+		branchRef := ref.Name()
+		branchName := strings.Split(string(branchRef), "refs/heads/")[1]
+
+		b := &config.Branch{
+			Name:  branchName,
+			Merge: branchRef,
+		}
+		if o.RemoteName == "" {
+			b.Remote = "origin"
+		} else {
+			b.Remote = o.RemoteName
+		}
+		if err := r.CreateBranch(b); err != nil {
+			return err
+		}
+	}
+
+	return nil
 }
 
 const (

repository_test.go

···
 	c.Assert(alt, IsNil)
 }
 
+func (s *RepositorySuite) TestCreateBranchAndBranch(c *C) {
+	r, _ := Init(memory.NewStorage(), nil)
+	testBranch := &config.Branch{
+		Name:   "foo",
+		Remote: "origin",
+		Merge:  "refs/heads/foo",
+	}
+	err := r.CreateBranch(testBranch)
+
+	c.Assert(err, IsNil)
+	cfg, err := r.Config()
+	c.Assert(err, IsNil)
+	c.Assert(len(cfg.Branches), Equals, 1)
+	branch := cfg.Branches["foo"]
+	c.Assert(branch.Name, Equals, testBranch.Name)
+	c.Assert(branch.Remote, Equals, testBranch.Remote)
+	c.Assert(branch.Merge, Equals, testBranch.Merge)
+
+	branch, err = r.Branch("foo")
+	c.Assert(err, IsNil)
+	c.Assert(branch.Name, Equals, testBranch.Name)
+	c.Assert(branch.Remote, Equals, testBranch.Remote)
+	c.Assert(branch.Merge, Equals, testBranch.Merge)
+}
+
+func (s *RepositorySuite) TestCreateBranchUnmarshal(c *C) {
+	r, _ := Init(memory.NewStorage(), nil)
+
+	expected := []byte(`[core]
+	bare = true
+[remote "foo"]
+	url = http://foo/foo.git
+	fetch = +refs/heads/*:refs/remotes/foo/*
+[branch "foo"]
+	remote = origin
+	merge = refs/heads/foo
+[branch "master"]
+	remote = origin
+	merge = refs/heads/master
+`)
+
+	_, err := r.CreateRemote(&config.RemoteConfig{
+		Name: "foo",
+		URLs: []string{"http://foo/foo.git"},
+	})
+	c.Assert(err, IsNil)
+	testBranch1 := &config.Branch{
+		Name:   "master",
+		Remote: "origin",
+		Merge:  "refs/heads/master",
+	}
+	testBranch2 := &config.Branch{
+		Name:   "foo",
+		Remote: "origin",
+		Merge:  "refs/heads/foo",
+	}
+	err = r.CreateBranch(testBranch1)
+	err = r.CreateBranch(testBranch2)
+
+	c.Assert(err, IsNil)
+	cfg, err := r.Config()
+	c.Assert(err, IsNil)
+	marshaled, err := cfg.Marshal()
+	c.Assert(string(expected), Equals, string(marshaled))
+}
+
+func (s *RepositorySuite) TestBranchInvalid(c *C) {
+	r, _ := Init(memory.NewStorage(), nil)
+	branch, err := r.Branch("foo")
+
+	c.Assert(err, NotNil)
+	c.Assert(branch, IsNil)
+}
+
+func (s *RepositorySuite) TestCreateBranchInvalid(c *C) {
+	r, _ := Init(memory.NewStorage(), nil)
+	err := r.CreateBranch(&config.Branch{})
+
+	c.Assert(err, NotNil)
+
+	testBranch := &config.Branch{
+		Name:   "foo",
+		Remote: "origin",
+		Merge:  "refs/heads/foo",
+	}
+	err = r.CreateBranch(testBranch)
+	c.Assert(err, IsNil)
+	err = r.CreateBranch(testBranch)
+	c.Assert(err, NotNil)
+}
+
+func (s *RepositorySuite) TestDeleteBranch(c *C) {
+	r, _ := Init(memory.NewStorage(), nil)
+	testBranch := &config.Branch{
+		Name:   "foo",
+		Remote: "origin",
+		Merge:  "refs/heads/foo",
+	}
+	err := r.CreateBranch(testBranch)
+
+	c.Assert(err, IsNil)
+
+	err = r.DeleteBranch("foo")
+	c.Assert(err, IsNil)
+
+	b, err := r.Branch("foo")
+	c.Assert(err, Equals, ErrBranchNotFound)
+	c.Assert(b, IsNil)
+
+	err = r.DeleteBranch("foo")
+	c.Assert(err, Equals, ErrBranchNotFound)
+}
+
 func (s *RepositorySuite) TestPlainInit(c *C) {
 	dir, err := ioutil.TempDir("", "plain-init")
 	c.Assert(err, IsNil)
···
 	remotes, err := r.Remotes()
 	c.Assert(err, IsNil)
 	c.Assert(remotes, HasLen, 1)
+	cfg, err := r.Config()
+	c.Assert(err, IsNil)
+	c.Assert(cfg.Branches, HasLen, 1)
+	c.Assert(cfg.Branches["master"].Name, Equals, "master")
 }
 
 func (s *RepositorySuite) TestPlainCloneContext(c *C) {
···
 	cfg, err := r.Config()
 	c.Assert(err, IsNil)
 	c.Assert(cfg.Remotes, HasLen, 1)
+	c.Assert(cfg.Branches, HasLen, 1)
 	c.Assert(cfg.Submodules, HasLen, 2)
 }
 
···
 	c.Assert(cfg.Remotes, HasLen, 1)
 	c.Assert(cfg.Remotes["origin"].Name, Equals, "origin")
 	c.Assert(cfg.Remotes["origin"].URLs, HasLen, 1)
+	c.Assert(cfg.Branches, HasLen, 1)
+	c.Assert(cfg.Branches["master"].Name, Equals, "master")
 }
 
 func (s *RepositorySuite) TestCloneSingleBranchAndNonHEAD(c *C) {
···
 	remotes, err := r.Remotes()
 	c.Assert(err, IsNil)
 	c.Assert(remotes, HasLen, 1)
+
+	cfg, err := r.Config()
+	c.Assert(err, IsNil)
+	c.Assert(cfg.Branches, HasLen, 1)
+	c.Assert(cfg.Branches["branch"].Name, Equals, "branch")
+	c.Assert(cfg.Branches["branch"].Remote, Equals, "origin")
+	c.Assert(cfg.Branches["branch"].Merge, Equals, plumbing.ReferenceName("refs/heads/branch"))
 
 	head, err = r.Reference(plumbing.HEAD, false)
 	c.Assert(err, IsNil)
···
 	c.Assert(err, IsNil)
 	c.Assert(remotes, HasLen, 1)
 
+	cfg, err := r.Config()
+	c.Assert(err, IsNil)
+	c.Assert(cfg.Branches, HasLen, 1)
+	c.Assert(cfg.Branches["master"].Name, Equals, "master")
+	c.Assert(cfg.Branches["master"].Remote, Equals, "origin")
+	c.Assert(cfg.Branches["master"].Merge, Equals, plumbing.ReferenceName("refs/heads/master"))
+
 	head, err = r.Reference(plumbing.HEAD, false)
 	c.Assert(err, IsNil)
 	c.Assert(head, NotNil)
···
 	})
 	c.Assert(err, IsNil)
 
+	cfg, err := r.Config()
+	c.Assert(err, IsNil)
+	c.Assert(cfg.Branches, HasLen, 0)
+
 	head, err := r.Reference(plumbing.HEAD, false)
 	c.Assert(err, IsNil)
 	c.Assert(head, NotNil)
···
 	})
 
 	c.Assert(err, IsNil)
+
+	cfg, err := r.Config()
+	c.Assert(err, IsNil)
+	c.Assert(cfg.Branches, HasLen, 0)
 
 	head, err := r.Reference(plumbing.HEAD, false)
 	c.Assert(err, IsNil)
···
 		ReferenceName: plumbing.ReferenceName("refs/tags/annotated-tag"),
 	})
 	c.Assert(err, IsNil)
+
+	cfg, err := r.Config()
+	c.Assert(err, IsNil)
+	c.Assert(cfg.Branches, HasLen, 0)
 
 	head, err := r.Reference(plumbing.HEAD, false)
 	c.Assert(err, IsNil)