<commit_before>package python

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"

	"strings"

	"sourcegraph.com/sourcegraph/srclib/graph"
	"sourcegraph.com/sourcegraph/srclib/grapher"
	"sourcegraph.com/sourcegraph/srclib/unit"
)

type GraphContext struct {
	Unit *unit.SourceUnit
	Reqs []*requirement
}

func NewGraphContext(unit *unit.SourceUnit) *GraphContext {
	var g GraphContext
	g.Unit = unit
	for _, dep := range unit.Dependencies {
		if req, err := asRequirement(dep); err == nil {
			g.Reqs = append(g.Reqs, req)
		}
	}
	return &g
}

// Graphs the Python source unit. If run outside of a Docker container, this assumes that the source unit has already
// been installed (via pip or `python setup.py install`).
func (c *GraphContext) Graph() (*grapher.Output, error) {
	if os.Getenv("IN_DOCKER_CONTAINER") != "" {
		// NOTE: this may cause an error when graphing any source unit that depends
		// on jedi (or any other dependency of the graph code)
		requirementFiles, err := filepath.Glob(filepath.Join(c.Unit.Dir, "*requirements.txt"))
		if err != nil {
			return nil, err
		}
		for _, requirementFile := range requirementFiles {
			exec.Command("pip", "install", "-r", requirementFile)
		}
		exec.Command("pip", "install", "-I", c.Unit.Dir)
	}

	cmd := exec.Command("python", "-m", "grapher.graph", c.Unit.Dir, "--verbose")
	cmd.Stderr = os.Stderr
	b, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	var raw RawOutput
	if err := json.Unmarshal(b, &raw); err != nil {
		return nil, err
	}

	out := c.transform(&raw, c.Unit)
	return out, nil
}

func (c *GraphContext) transform(raw *RawOutput, unit *unit.SourceUnit) *grapher.Output {
	var out grapher.Output

	for _, def := range raw.Defs {
		out.Defs = append(out.Defs, c.transformDef(def))
		if doc := c.transformDefDoc(def); doc != nil {
			out.Docs = append(out.Docs, doc)
		}
	}
	for _, ref := range raw.Refs {
		if outRef, err := c.transformRef(ref); err == nil {
			out.Refs = append(out.Refs, outRef)
		} else {
			log.Printf("Could not transform ref %v: %s", ref, err)
		}
	}

	return &out
}

var jediKindToDefKind = map[string]graph.DefKind{
	"statement":        graph.Var,
	"statementelement": graph.Var,
	"param":            graph.Var,
	"module":           graph.Module,
	"submodule":        graph.Module,
	"class":            graph.Type,
	"function":         graph.Func,
	"lambda":           graph.Func,
	"import":           graph.Var,
}

func (c *GraphContext) transformDef(rawDef *RawDef) *graph.Def {
	return &graph.Def{
		DefKey: graph.DefKey{
			Repo:     c.Unit.Repo,
			Unit:     c.Unit.Name,
			UnitType: c.Unit.Type,
			Path:     graph.DefPath(rawDef.Path),
		},
		TreePath: graph.TreePath(rawDef.Path), // TODO: make this consistent w/ old way
		Kind:     jediKindToDefKind[rawDef.Kind],
		Name:     rawDef.Name,
		File:     rawDef.File,
		DefStart: rawDef.DefStart,
		DefEnd:   rawDef.DefEnd,
		Exported: rawDef.Exported,
		Data:     nil, // TODO
	}
}

func (c *GraphContext) transformRef(rawRef *RawRef) (*graph.Ref, error) {
	defUnit, err := c.inferSourceUnit(rawRef, c.Reqs)
	if err != nil {
		return nil, err
	}

	defPath := graph.DefPath(rawRef.DefPath)
	if defPath == "" {
		defPath = "."
	}

	return &graph.Ref{
		DefRepo:     defUnit.Repo,
		DefUnitType: defUnit.Type,
		DefUnit:     defUnit.Name,
		DefPath:     defPath,

		Repo:     c.Unit.Repo,
		Unit:     c.Unit.Name,
		UnitType: c.Unit.Type,

		File:  rawRef.File,
		Start: rawRef.Start,
		End:   rawRef.End,
	}, nil
}

func (c *GraphContext) transformDefDoc(rawDef *RawDef) *graph.Doc {
	return nil
}

func (c *GraphContext) inferSourceUnit(rawRef *RawRef, reqs []*requirement) (*unit.SourceUnit, error) {
	if rawRef.ToBuiltin {
		return stdLibPkg.SourceUnit(), nil
	}
	return c.inferSourceUnitFromFile(rawRef.DefFile, reqs)
}

// Note: file is expected to be an absolute path
func (c *GraphContext) inferSourceUnitFromFile(file string, reqs []*requirement) (*unit.SourceUnit, error) {
	// Case: in current source unit (u)
	pwd, _ := os.Getwd()
	if isSubPath(pwd, file) {
		return c.Unit, nil
	}

	// Case: in dependent source unit(depUnits)
	fileCmps := strings.Split(file, string(filepath.Separator))
	pkgsDirIdx := -1
	for i, cmp := range fileCmps {
		if cmp == "site-packages" || cmp == "dist-packages" {
			pkgsDirIdx = i
			break<commit_msg>emit whether ref is a defref<commit_after>package python

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"

	"strings"

	"sourcegraph.com/sourcegraph/srclib/graph"
	"sourcegraph.com/sourcegraph/srclib/grapher"
	"sourcegraph.com/sourcegraph/srclib/unit"
)

type GraphContext struct {
	Unit *unit.SourceUnit
	Reqs []*requirement
}

func NewGraphContext(unit *unit.SourceUnit) *GraphContext {
	var g GraphContext
	g.Unit = unit
	for _, dep := range unit.Dependencies {
		if req, err := asRequirement(dep); err == nil {
			g.Reqs = append(g.Reqs, req)
		}
	}
	return &g
}

// Graphs the Python source unit. If run outside of a Docker container, this assumes that the source unit has already
// been installed (via pip or `python setup.py install`).
func (c *GraphContext) Graph() (*grapher.Output, error) {
	if os.Getenv("IN_DOCKER_CONTAINER") != "" {
		// NOTE: this may cause an error when graphing any source unit that depends
		// on jedi (or any other dependency of the graph code)
		requirementFiles, err := filepath.Glob(filepath.Join(c.Unit.Dir, "*requirements.txt"))
		if err != nil {
			return nil, err
		}
		for _, requirementFile := range requirementFiles {
			exec.Command("pip", "install", "-r", requirementFile)
		}
		exec.Command("pip", "install", "-I", c.Unit.Dir)
	}

	cmd := exec.Command("python", "-m", "grapher.graph", c.Unit.Dir, "--verbose")
	cmd.Stderr = os.Stderr
	b, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	var raw RawOutput
	if err := json.Unmarshal(b, &raw); err != nil {
		return nil, err
	}

	out := c.transform(&raw, c.Unit)
	return out, nil
}

func (c *GraphContext) transform(raw *RawOutput, unit *unit.SourceUnit) *grapher.Output {
	var out grapher.Output

	for _, def := range raw.Defs {
		out.Defs = append(out.Defs, c.transformDef(def))
		if doc := c.transformDefDoc(def); doc != nil {
			out.Docs = append(out.Docs, doc)
		}
	}
	for _, ref := range raw.Refs {
		if outRef, err := c.transformRef(ref); err == nil {
			out.Refs = append(out.Refs, outRef)
		} else {
			log.Printf("Could not transform ref %v: %s", ref, err)
		}
	}

	return &out
}

var jediKindToDefKind = map[string]graph.DefKind{
	"statement":        graph.Var,
	"statementelement": graph.Var,
	"param":            graph.Var,
	"module":           graph.Module,
	"submodule":        graph.Module,
	"class":            graph.Type,
	"function":         graph.Func,
	"lambda":           graph.Func,
	"import":           graph.Var,
}

func (c *GraphContext) transformDef(rawDef *RawDef) *graph.Def {
	return &graph.Def{
		DefKey: graph.DefKey{
			Repo:     c.Unit.Repo,
			Unit:     c.Unit.Name,
			UnitType: c.Unit.Type,
			Path:     graph.DefPath(rawDef.Path),
		},
		TreePath: graph.TreePath(rawDef.Path), // TODO: make this consistent w/ old way
		Kind:     jediKindToDefKind[rawDef.Kind],
		Name:     rawDef.Name,
		File:     rawDef.File,
		DefStart: rawDef.DefStart,
		DefEnd:   rawDef.DefEnd,
		Exported: rawDef.Exported,
		Data:     nil, // TODO
	}
}

func (c *GraphContext) transformRef(rawRef *RawRef) (*graph.Ref, error) {
	defUnit, err := c.inferSourceUnit(rawRef, c.Reqs)
	if err != nil {
		return nil, err
	}

	defPath := graph.DefPath(rawRef.DefPath)
	if defPath == "" {
		defPath = "."
	}

	return &graph.Ref{
		DefRepo:     defUnit.Repo,
		DefUnitType: defUnit.Type,
		DefUnit:     defUnit.Name,
		DefPath:     defPath,

		Repo:     c.Unit.Repo,
		Unit:     c.Unit.Name,
		UnitType: c.Unit.Type,

		File:  rawRef.File,
		Start: rawRef.Start,
		End:   rawRef.End,
		Def:   rawRef.Def,
	}, nil
}

func (c *GraphContext) transformDefDoc(rawDef *RawDef) *graph.Doc {
	return nil
}

func (c *GraphContext) inferSourceUnit(rawRef *RawRef, reqs []*requirement) (*unit.SourceUnit, error) {
	if rawRef.ToBuiltin {
		return stdLibPkg.SourceUnit(), nil
	}
	return c.inferSourceUnitFromFile(rawRef.DefFile, reqs)
}

// Note: file is expected to be an absolute path
func (c *GraphContext) inferSourceUnitFromFile(file string, reqs []*requirement) (*unit.SourceUnit, error) {
	// Case: in current source unit (u)
	pwd, _ := os.Getwd()
	if isSubPath(pwd, file) {
		return c.Unit, nil
	}

	// Case: in dependent source unit(depUnits)
	fileCmps := strings.Split(file, string(filepath.Separator))
	pkgsDirIdx := -1
	for i, cmp := range fileCmps {
		if cmp == "site-packages" || cmp == "dist-packages" {
			pkgsDirIdx = i
			break
{'fraction_non_alphanumeric': 0.1442528735632184, 'fraction_numerical': 0.00022988505747126436, 'mean_word_length': 3.3483258370814593, 'pattern_counts': {'":': 18, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6768664', 'n_tokens_mistral': 3356, 'n_tokens_neox': 3196, 'n_words': 802}