X-Git-Url: http://nitlanguage.org

diff --git a/src/nituml.nit b/src/nituml.nit
index 3d5f3ba..2723a8b 100644
--- a/src/nituml.nit
+++ b/src/nituml.nit
@@ -20,15 +20,38 @@ import frontend
 import uml
 
 redef class ToolContext
+
+	# Phase that generates UML diagrams from model entities.
 	var umlphase: Phase = new UMLPhase(self, null)
+
+	# What to generate?
+	var opt_gen = new OptionEnum(["class", "package"], "Choose which type of uml diagram to generate", 0, "--diagram")
+
+	# Generate private?
+	var opt_privacy = new OptionBool("Generates private API", "-p", "--private")
+
+	redef init do
+		option_context.add_option opt_gen
+		option_context.add_option opt_privacy
+		super
+	end
 end
 
 private class UMLPhase
 	super Phase
 	redef fun process_mainmodule(mainmodule, mmodules)
 	do
-		var d = new UMLModel(mainmodule.model, mainmodule, toolcontext)
-		print d.generate_class_uml.write_to_string
+		var filters = new ModelFilter
+		if not toolcontext.opt_privacy.value then
+			filters.min_visibility = protected_visibility
+		end
+
+		var d = new UMLModel(toolcontext.modelbuilder.model, mainmodule, filters)
+		if toolcontext.opt_gen.value == 0 then
+			print d.generate_class_uml.write_to_string
+		else if toolcontext.opt_gen.value == 1 then
+			print d.generate_package_uml.write_to_string
+		end
 	end
 end
 
@@ -40,7 +63,7 @@ var arguments = toolcontext.option_context.rest
 # build model
 var model = new Model
 var mbuilder = new ModelBuilder(model, toolcontext)
-var mmodules = mbuilder.parse(arguments)
+var mmodules = mbuilder.parse_full(arguments)
 
 if mmodules.is_empty then return
 mbuilder.run_phases