Runner.AddInput will no longer throw NullReferenceException (migueldeicaza#418)

csteegz · migueldeicaza · commit a8a280e31ff5 · 2019-10-25T15:46:11.000-04:00
diff --git a/TensorFlowSharp/Tensorflow.cs b/TensorFlowSharp/Tensorflow.cs
@@ -2929,14 +2929,28 @@ public Runner AddTarget (params TFOperation [] targets)
 			// Parses user strings that contain both the operation name and an index.
 			TFOutput ParseOutput (string operation)
 			{
+				TFOperation tfOperation;
 				var p = operation.IndexOf (':');
-				if (p != -1 && p != operation.Length - 1){
+				if (p != -1 && p != operation.Length - 1)
+				{
+					
 					var op = operation.Substring (0, p);
-					if (int.TryParse (operation.Substring (p + 1), out var idx)){
-						return session.Graph [op] [idx];
+					if (int.TryParse (operation.Substring (p + 1), out var idx))
+					{
+						tfOperation = session.Graph [op];
+						if (tfOperation == null)
+						{
+							throw new ArgumentOutOfRangeException($"An operation named {op} is not in this graph.");
+						}
+						return tfOperation [idx];
 					}
 				}
-				return session.Graph [operation] [0];
+				tfOperation = session.Graph [operation];
+				if (tfOperation == null)
+				{
+					throw new ArgumentOutOfRangeException($"An operation named {operation} is not in this graph.");
+				}
+				return tfOperation [0];
 			}
 
 			/// <summary>
diff --git a/tests/TensorFlowSharp.Tests.CSharp/SessionTests.cs b/tests/TensorFlowSharp.Tests.CSharp/SessionTests.cs
@@ -1,4 +1,5 @@
-&#65279;using System.Linq;
+&#65279;using System;
+using System.Linq;
 using TensorFlow;
 using Xunit;
 
@@ -16,5 +17,18 @@ public void Should_ListDeviceReturnDevices ()
 				Assert.True(devices.Any());
 			}
 		}
+		
+		[Theory]
+		[InlineData("Placeholder")]
+		[InlineData("Placeholder:0")]
+		public void ParseOutput_ThrowsForMissingOp (string name)
+		{
+			using (var graph = new TFGraph ())
+			using (var session = new TFSession (graph))
+			{
+				var runner = session.GetRunner();
+				Assert.Throws<ArgumentOutOfRangeException>(() => runner.AddInput(name, new TFTensor(1)));
+			}
+		}
 	}
 }

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`		`-using System.Linq;`
	`1`	`+using System;`
	`2`	`+using System.Linq;`
`2`	`3`	`using TensorFlow;`
`3`	`4`	`using Xunit;`
`4`	`5`
`@@ -16,5 +17,18 @@ public void Should_ListDeviceReturnDevices ()`
`16`	`17`	`Assert.True(devices.Any());`
`17`	`18`	`}`
`18`	`19`	`}`
	`20`	`+`
	`21`	`+ [Theory]`
	`22`	`+ [InlineData("Placeholder")]`
	`23`	`+ [InlineData("Placeholder:0")]`
	`24`	`+ public void ParseOutput_ThrowsForMissingOp (string name)`
	`25`	`+ {`
	`26`	`+ using (var graph = new TFGraph ())`
	`27`	`+ using (var session = new TFSession (graph))`
	`28`	`+ {`
	`29`	`+ var runner = session.GetRunner();`
	`30`	`+ Assert.Throws<ArgumentOutOfRangeException>(() => runner.AddInput(name, new TFTensor(1)));`
	`31`	`+ }`
	`32`	`+ }`
`19`	`33`	`}`
`20`	`34`	`}`